0713-1718

diffusion.py

@@ -619,7 +619,7 @@ def generate_samples(ddpm21cm, num_new_img, max_num_img_per_gpu, rank, world_siz
     samples_list = [np.empty_like(samples) for _ in range(world_size)]
     dist.all_gather_object(samples_list, samples)
 
-    if rank == 1:
+    if rank == 0:
         all_samples = np.concatenate(samples_list, axis=0)
         return all_samples
     else:
@@ -631,12 +631,10 @@ def sample(rank, world_size, config, num_new_img, max_num_img_per_gpu, return_di
 
     samples = generate_samples(ddpm21cm, num_new_img, max_num_img_per_gpu, rank, world_size)
 
-    print(f"device {torch.cuda.current_device()}, rank = {rank}, keys = {return_dict.keys()}, samples.shape = {np.shape(samples)}")
-
-    if rank == 1:
+    # print(f"device {torch.cuda.current_device()}, rank = {rank}, keys = {return_dict.keys()}, samples.shape = {np.shape(samples)}")
+    if rank == 0:
         return_dict['samples'] = samples
-
-    print(f"device {torch.cuda.current_device()}, rank = {rank}, keys = {return_dict.keys()}")
+    # print(f"device {torch.cuda.current_device()}, rank = {rank}, keys = {return_dict.keys()}")
 
     dist.destroy_process_group()
 
@@ -664,8 +662,8 @@ if __name__ == "__main__":
 
         mp.spawn(sample, args=(world_size, config, num_new_img, max_num_img_per_gpu, return_dict), nprocs=world_size, join=True)
 
-        print("---"*30)
-        print(f"device {torch.cuda.current_device()}, keys = {return_dict.keys()}")
+        # print("---"*30)
+        # print(f"device {torch.cuda.current_device()}, keys = {return_dict.keys()}")
         if "samples" in return_dict:
             samples = return_dict["samples"]
             print(f"device {torch.cuda.current_device()} generated samples shape: {samples.shape}")