default world size

LAION-AI · Feb 26, 2023 · cc77d27 · cc77d27
1 parent e94e289
commit cc77d27
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/model/model_training/trainer_sft.py b/model/model_training/trainer_sft.py
@@ -95,7 +95,6 @@ def prediction_step(
         prediction_loss_only: bool,
         ignore_keys: Optional[List[str]] = None,
     ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
-
         with torch.no_grad():
             loss, logits, labels, labels_mask = self._compute_loss(model, inputs)
             labels[~labels_mask.bool()] = -100  # padding_index
@@ -190,7 +189,7 @@ def argument_parsing(notebook=False, notebook_args=None):
 
     # get the world size in deeepspeed
     if conf["deepspeed"]:
-        conf["world_size"] = int(os.environ["WORLD_SIZE"])
+        conf["world_size"] = int(os.getenv("WORLD_SIZE", default="1"))
     else:
         conf["world_size"] = 1