fix deepspeed model saving (#370)

huggingface · May 18, 2022 · d33dc39 · d33dc39
1 parent 043d2ec
commit d33dc39
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -881,7 +881,7 @@ def get_state_dict(self, model):
                 is_zero_3 = self.state.deepspeed_plugin.zero_stage == 3
 
         if is_zero_3:
-            state_dict = model._zero3_consolidated_fp16_state_dict()
+            state_dict = model._zero3_consolidated_16bit_state_dict()
         else:
             model = self.unwrap_model(model)
             state_dict = model.state_dict()