0710-1102

diffusion.py

@@ -71,6 +71,7 @@ def ddp_setup(rank: int, world_size: int):
   """
   os.environ["MASTER_ADDR"] = "localhost"
   os.environ["MASTER_PORT"] = "12355"
+  print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!ddp_setup, rank =", rank)
   torch.cuda.set_device(rank)
   init_process_group(backend="nccl", rank=rank, world_size=world_size)
 
@@ -237,7 +238,7 @@ class TrainConfig:
     stride = (2,2) if dim == 2 else (2,2,1)
     num_image = 2000#32000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
     batch_size = 2#2#50#20#2#100 # 10
-    n_epoch = 10#50#20#20#2#5#25 # 120
+    n_epoch = 2#10#50#20#20#2#5#25 # 120
     HII_DIM = 28#64
     num_redshift = 4#128#64#512#256#256#64#512#128
     channel = 1