diff --git a/stoke/status.py b/stoke/status.py index 7b03f59..55c12c6 100644 --- 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 if self._configs.get("DeepspeedConfig") else None, + "zero": self._configs.get("DeepspeedConfig").zero_optimization.stage + if self._configs.get("DeepspeedConfig") + else None, "oss": fairscale_oss, "sharded": fairscale_sddp, "world_size": -1,