0710-1805

diffusion.py

@@ -259,7 +259,7 @@ class TrainConfig:
     n_param = 2
     guide_w = 0#-1#0#-1#0#-1#0.1#[0,0.1] #[0,0.5,2] strength of generative guidance
     drop_prob = 0#0.28 # only takes effect when guide_w != -1
-    ema=True # whether to use ema
+    ema=False # whether to use ema
     ema_rate=0.995
 
     # seed = 0
@@ -482,8 +482,8 @@ class DDPM21CM:
                     if self.config.save_name:
                         model_state = {
                             'epoch': ep,
-                            'unet_state_dict': self.nn_model.state_dict(),
-                            'ema_unet_state_dict': self.ema_model.state_dict(),
+                            'unet_state_dict': self.nn_model.module.state_dict(),
+                            # 'ema_unet_state_dict': self.ema_model.state_dict(),
                             }
                         torch.save(model_state, self.config.save_name+f"-N{self.config.num_image}")
                         print('saved model at ' + self.config.save_name+f"-N{self.config.num_image}")