Skip to content

Commit

Permalink
Minor reformatting
Browse files Browse the repository at this point in the history
  • Loading branch information
RandomDefaultUser committed Apr 26, 2024
1 parent 2cde3f1 commit bd68063
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions mala/common/parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -1335,15 +1335,18 @@ def use_ddp(self, value):
if value:
print("initializing torch.distributed.")
# JOSHR:
# We start up torch distributed here. As is fairly standard convention, we get the rank
# and world size arguments via environment variables (RANK, WORLD_SIZE). In addition to
# those variables, LOCAL_RANK, MASTER_ADDR and MASTER_PORT should be set.
# We start up torch distributed here. As is fairly standard
# convention, we get the rank and world size arguments via
# environment variables (RANK, WORLD_SIZE). In addition to
# those variables, LOCAL_RANK, MASTER_ADDR and MASTER_PORT
# should be set.
rank = int(os.environ.get("RANK"))
world_size = int(os.environ.get("WORLD_SIZE"))

dist.init_process_group("nccl", rank=rank, world_size=world_size)

# Invalidate, will be updated in setter.
set_ddp_status(value)
# Invalidate, will be updated in setter.
self.device = None
self._use_ddp = value
self.network._update_ddp(self.use_ddp)
Expand Down

0 comments on commit bd68063

Please sign in to comment.