diffusion_models.diffusion_trainer
==================================

.. py:module:: diffusion_models.diffusion_trainer




Module Contents
---------------

.. py:class:: DiffusionTrainer(model, dataset, optimizer, training_configuration, loss_function = F.l1_loss, scheduler = None, log_configuration = LogConfiguration(), reverse_transforms = lambda x: x, device = 'cuda')

   .. py:attribute:: model


   .. py:attribute:: optimizer


   .. py:attribute:: loss_function


   .. py:attribute:: training_configuration


   .. py:attribute:: scheduler


   .. py:attribute:: device


   .. py:attribute:: dataloader


   .. py:attribute:: scaler


   .. py:attribute:: log_configuration


   .. py:attribute:: checkpoint_path


   .. py:attribute:: tensorboard_manager


   .. py:attribute:: reverse_transforms


   .. py:attribute:: benchmark
      :value: True



   .. py:method:: save_checkpoint(epoch, checkpoint_name)


   .. py:method:: train()


   .. py:method:: log_to_tensorboard(metrics, global_step)