change dist to torch.distributed to fix bug in assert. (#638)

microsoft · Jan 5, 2021 · d38ad6a · d38ad6a
1 parent a9a83a6
commit d38ad6a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/deepspeed/utils/distributed.py b/deepspeed/utils/distributed.py
@@ -77,9 +77,9 @@ def mpi_discovery(distributed_port=TORCH_DISTRIBUTED_DEFAULT_PORT, verbose=True)
                     os.environ['MASTER_PORT']))
 
     if torch.distributed.is_initialized():
-        assert dist.get_rank() == rank, "MPI rank {} does not match torch rank {}".format(rank, dist.get_rank())
-        assert dist.get_world_size() == world_size, "MPI world size {} does not match torch world size {}".format(
-            world_size, dist.get_world_size())
+        assert torch.distributed.get_rank() == rank, "MPI rank {} does not match torch rank {}".format(rank, dist.get_rank())
+        assert torch.distributed.get_world_size() == world_size, "MPI world size {} does not match torch world size {}".format(
+            world_size, torch.distributed.get_world_size())
 
 
 def in_aml():