fix bugs when training with single device

tqch · Sep 14, 2022 · 9765b3a · 9765b3a
1 parent a0e0d39
commit 9765b3a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -11,7 +11,7 @@
 
 
 def logger(msg, **kwargs):
- if dist.get_rank() == 0:
+ if dist.is_initialized() and dist.get_rank() == 0:
  print(msg, **kwargs)
 
 
@@ -72,6 +72,7 @@ def main(args):
  model = DDP(_model, device_ids=[rank, ])
  train_device = torch.device(f"cuda:{rank}")
  else:
+ rank = 0
  model = _model.to(train_device)
 
  optimizer = Adam(model.parameters(), lr=lr, betas=(beta1, beta2))