0710-1820

diffusion.py

@@ -556,7 +556,7 @@ class DDPM21CM:
 # print("device =", config.device)
 
 # %%
-def single_main(rank, world_size):
+def main(rank, world_size):
     config = TrainConfig()
     ddp_setup(rank, world_size)
     
@@ -569,6 +569,7 @@ def single_main(rank, world_size):
         print(f" num_image = {ddpm21cm.config.num_image} ".center(50, '-'))
         print(f"run_name = {ddpm21cm.config.run_name}")
         ddpm21cm.train()
+        destroy_process_group()
 
         
 if __name__ == "__main__":
@@ -576,7 +577,7 @@ if __name__ == "__main__":
     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
     world_size = 2#torch.cuda.device_count()
 
-    mp.spawn(single_main, args=(world_size,), nprocs=world_size)
+    mp.spawn(main, args=(world_size,), nprocs=world_size)
     # notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')
 
 # %%