Use gloo backend on Windows for Pytorch

This commit is contained in:
Justin John
2023-02-27 10:36:35 +05:30
committed by GitHub
parent bf7042e454
commit 730365fcde
+3 -2
View File
@@ -65,7 +65,8 @@ def run(rank, n_gpus, hps):
writer = SummaryWriter(log_dir=hps.model_dir)
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
# Use gloo backend on Windows for Pytorch
dist.init_process_group(backend= 'gloo' if os.name == 'nt' else 'nccl', init_method='env://', world_size=n_gpus, rank=rank)
torch.manual_seed(hps.train.seed)
torch.cuda.set_device(rank)
@@ -317,4 +318,4 @@ def evaluate(hps, generator, eval_loader, writer_eval):
if __name__ == "__main__":
main()
main()