Xsmos
/

ml21cm

generate 21cm lightcones

denoising diffusion probabilistic model

Model card Files Files and versions

Metrics Training metrics Community

Xsmos commited on Jul 11, 2024

Commit

9334565

·

verified ·

1 Parent(s): c194d7c

0710-2302

Files changed (1) hide show

diffusion.py +1 -1

diffusion.py CHANGED Viewed

@@ -575,7 +575,7 @@ def main(rank, world_size):
 if __name__ == "__main__":
     # torch.multiprocessing.set_start_method("spawn")
     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
-    world_size = 1#torch.cuda.device_count()
     mp.spawn(main, args=(world_size,), nprocs=world_size)
     # notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')

 if __name__ == "__main__":
     # torch.multiprocessing.set_start_method("spawn")
     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
+    world_size = 2#torch.cuda.device_count()
     mp.spawn(main, args=(world_size,), nprocs=world_size)
     # notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')