Xsmos
/

ml21cm

generate 21cm lightcones

denoising diffusion probabilistic model

Model card Files Files and versions

Metrics Training metrics Community

Xsmos commited on Jul 13, 2024

Commit

3ceda46

·

verified ·

1 Parent(s): 1ae487f

0713-1459

Files changed (1) hide show

diffusion.py +1 -1

diffusion.py CHANGED Viewed

@@ -532,7 +532,7 @@ class DDPM21CM:
             self.nn_model.module.load_state_dict(torch.load(file)['ema_unet_state_dict'])
         else:
             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()
         # self.nn_model.to(self.ddpm.device)

             self.nn_model.module.load_state_dict(torch.load(file)['ema_unet_state_dict'])
         else:
             self.nn_model.module.load_state_dict(torch.load(file)['unet_state_dict'])
+        print(f"device {torch.cuda.current_device()} resumed nn_model from {file}")
         # nn_model = ContextUnet(n_param=1, image_size=28)
         # nn_model.train()
         # self.nn_model.to(self.ddpm.device)