Update engine.py to avoid torch warning (microsoft#5408)

The state_dict function of module.py from torch write a warning if arguments are positional arguments and not keyword arguments --------- Co-authored-by: Logan Adams <[email protected]> Co-authored-by: ebonnafoux <[email protected]>
dbyoung18 · Jun 11, 2024 · 4b9bae8 · 4b9bae8
1 parent ace7cc7
commit 4b9bae8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -2542,7 +2542,7 @@ def all_gather_scalar(self, value, dp_group):
         return tensor_list
 
     def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclude_frozen_parameters=False):
-        sd = self.module.state_dict(destination, prefix, keep_vars)
+        sd = self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
 
         # Remove frozen parameter weights from state_dict if specified
         if exclude_frozen_parameters: