0712-1558

diffusion.py

@@ -528,9 +528,9 @@ class DDPM21CM:
 
         # nn_model = ContextUnet(n_param=self.config.n_param, image_size=self.config.HII_DIM, dim=self.config.dim, stride=self.config.stride).to(self.config.device)
         if ema:
-            self.nn_model.load_state_dict(torch.load(file)['ema_unet_state_dict'])
+            self.nn_model.module.load_state_dict(torch.load(file)['ema_unet_state_dict'])
         else:
-            self.nn_model.load_state_dict(torch.load(file)['unet_state_dict'])
+            self.nn_model.module.load_state_dict(torch.load(file)['unet_state_dict'])
         print(f"nn_model resumed from {file}")
         # nn_model = ContextUnet(n_param=1, image_size=28)
         # nn_model.train()
@@ -601,15 +601,15 @@ if __name__ == "__main__":
 
 # %%
 
-def generate_samples(model, num_new_img, max_num_img_per_gpu, rank, world_size):
+def generate_samples(ddpm21cm, num_new_img, max_num_img_per_gpu, rank, world_size):
     samples = []
     for _ in range(num_new_img // max_num_img_per_gpu):    
-        sample = model.module.sample(filename, params=torch.tensor([4.4, 131.341]), num_new_img=max_num_img_per_gpu)
+        sample = ddpm21cm.sample(filename, params=torch.tensor([4.4, 131.341]), num_new_img=max_num_img_per_gpu)
         samples.append(sample)
-        # model.sample(filename, params=torch.tensor((5.6, 19.037)), num_new_img=max_num_img_per_gpu)
-        # model.sample(filename, params=torch.tensor((4.699, 30)), num_new_img=max_num_img_per_gpu)
-        # model.sample(filename, params=torch.tensor((5.477, 200)), num_new_img=max_num_img_per_gpu)
-        # model.sample(filename, params=torch.tensor((4.8, 131.341)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(filename, params=torch.tensor((5.6, 19.037)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(filename, params=torch.tensor((4.699, 30)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(filename, params=torch.tensor((5.477, 200)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(filename, params=torch.tensor((4.8, 131.341)), num_new_img=max_num_img_per_gpu)
     samples = np.concatenate(samples, axis=0)
 
     samples_list = [np.empty_like(samples) for _ in range(world_size)]