deepspeed save model temp fix (#374)

* fix deepspeed model saving * fix deepspeed zero stage-3 model save fixes #369 Co-Authored-By: Kovvuri Satyanarayana Reddy <54667784+KOVVURISATYANARAYANAREDDY@users.noreply.github.com> Co-authored-by: Kovvuri Satyanarayana Reddy <54667784+KOVVURISATYANARAYANAREDDY@users.noreply.github.com>
huggingface · May 19, 2022 · 6163e20 · 6163e20
1 parent d33dc39
commit 6163e20
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -886,9 +886,10 @@ def get_state_dict(self, model):
             model = self.unwrap_model(model)
             state_dict = model.state_dict()
 
-        for k in state_dict:
-            if state_dict[k].dtype == torch.float16:
-                state_dict[k] = state_dict[k].float()
+        if state_dict is not None:
+            for k in state_dict:
+                if state_dict[k].dtype == torch.float16:
+                    state_dict[k] = state_dict[k].float()
 
         return state_dict