Update app.py

app.py

@@ -101,7 +101,7 @@ class UNetWrapper:
 
     def save_checkpoint(self, save_path):
         """Save checkpoint with model, optimizer, and scheduler states."""
-        save_dict = {
+        self.save_dict = {
             'model_state_dict': self.model.state_dict(),
             'optimizer_state_dict': self.optimizer.state_dict(),
             'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
@@ -112,7 +112,7 @@ class UNetWrapper:
             'epoch': self.epoch,
             'loss': self.loss
         }
-        torch.save(save_dict, save_path)
+        torch.save(self.save_dict, save_path)
         print(f"Checkpoint saved at epoch {self.epoch}, loss: {self.loss}")
 
     def load_checkpoint(self, checkpoint_path):
@@ -162,8 +162,8 @@ pipeline_tag: image-to-image
 
 ## Model Description
 Custom UNet model for Pix2Pix image translation.
-- **Image Size:** {save_dict['model_config']['img_size']}
-- **Model Type:** {"big" if big else "small"}_UNet ({save_dict['model_config']['img_size']})
+- **Image Size:** {self.save_dict['model_config']['img_size']}
+- **Model Type:** {"big" if big else "small"}_UNet ({self.save_dict['model_config']['img_size']})
 
 ## Usage
 
@@ -188,8 +188,8 @@ model.eval()
         # Save and upload README
         with open("README.md", "w") as f:
             f.write(f"# Pix2Pix UNet Model\n\n"
-                    f"- **Image Size:** {save_dict['model_config']['img_size']}\n"
-                    f"- **Model Type:** {'big' if big else 'small'}_UNet ({save_dict['model_config']['img_size']})\n"
+                    f"- **Image Size:** {self.save_dict['model_config']['img_size']}\n"
+                    f"- **Model Type:** {'big' if big else 'small'}_UNet ({self.save_dict['model_config']['img_size']})\n"
                     f"## Model Architecture\n{str(self.model)}")
         
         self.api.upload_file(