Merge pull request #8 from fidelity/fix

bug with zero stage config print
fidelity · Jul 28, 2021 · bb2142c · bb2142c
2 parents 3ced64d + 6061903
commit bb2142c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/stoke/status.py b/stoke/status.py
@@ -162,7 +162,9 @@ def __init__(
             "grad_clip": grad_clip,
             "gpu": gpu,
             "distributed": distributed,
-            "zero": self._configs.get("DeepspeedConfig").zero_optimization.stage,
+            "zero": self._configs.get("DeepspeedConfig").zero_optimization.stage
+            if self._configs.get("DeepspeedConfig")
+            else None,
             "oss": fairscale_oss,
             "sharded": fairscale_sddp,
             "world_size": -1,