Saving whole model instead of weights only

src/trainers/diffusion_trainer.py

@@ -236,7 +236,7 @@ class DiffusionTrainer:
             )
 
     def save_checkpoint(self, val_loss, task, iteration: int):
-        torch.save(self.model.state_dict(), "checkpoint.pt")
+        torch.save(self.model, "checkpoint.pt")
         task.update_output_model(
             model_path="checkpoint.pt", iteration=iteration, auto_delete_file=False
         )

src/trainers/trainer.py

@@ -279,7 +279,7 @@ class Trainer:
         return test_loss
 
     def save_checkpoint(self, val_loss, task, iteration: int):
-        torch.save(self.model.state_dict(), "checkpoint.pt")
+        torch.save(self.model, "checkpoint.pt")
         task.update_output_model(
             model_path="checkpoint.pt", iteration=iteration, auto_delete_file=False
         )