0715-1132

diffusion.py

@@ -336,8 +336,8 @@ class DDPM21CM:
             self.nn_model.module.load_state_dict(torch.load(config.resume)['unet_state_dict'])
             print(f"device {torch.cuda.current_device()} resumed nn_model from {config.resume}")
 
-        # self.number_of_params = sum(x.numel() for x in self.nn_model.parameters())
-        # print(f"Number of parameters for nn_model: {self.number_of_params}")
+        self.number_of_params = sum(x.numel() for x in self.nn_model.parameters())
+        print(f" Number of parameters for nn_model: {self.number_of_params} ".center(100,'-'))
 
         # whether to use ema
         if config.ema: