0715-1314

diffusion.py

@@ -282,7 +282,7 @@ class TrainConfig:
     # params =  params
     # data_dir = './data' # data directory
 
-    use_fp16 = False
+    use_fp16 = True
     dtype = torch.float16 if use_fp16 else torch.float32
     mixed_precision = "fp16"
     gradient_accumulation_steps = 1
@@ -437,7 +437,7 @@ class DDPM21CM:
                 with self.accelerator.accumulate(self.nn_model):
                     x = x.to(self.config.device)
                     print("x = x.to(self.config.device), x.dtype =", x.dtype)
-                    x = x.to(self.config.dtype)
+                    # x = x.to(self.config.dtype)
                     print("x = x.to(self.dtype), x.dtype =", x.dtype)
                     xt, noise, ts = self.ddpm.add_noise(x)