0715-1359

diffusion.py

@@ -589,10 +589,10 @@ def train(rank, world_size):
 
         
 if __name__ == "__main__":
-    print(" training ".center(100,'-'))
+    world_size = torch.cuda.device_count()
+    print(f" training, world_size = {world_size} ".center(100,'-'))
     # torch.multiprocessing.set_start_method("spawn")
     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
-    world_size = torch.cuda.device_count()
 
     mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
     # notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')