Fix typo

microsoft · Aug 13, 2024 · bf84522 · bf84522
1 parent 348b29d
commit bf84522
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
@@ -447,7 +447,7 @@ def wrapper(module, *args, **kwargs):
                     # since skip_init won't involve any computations or weight adjustments, we can directly utilize post_init
                     self._post_init_method(_module)
                     return _module
-                wrapped._ds_has_wrapped = True
+                wrapper._ds_has_wrapped = True
                 return wrapper
 
             def post_wrapper_to_empty(f):
@@ -520,7 +520,7 @@ def wrapper(module, *args, **kwargs):
                 print_rank_0(f'After initializing followed by post init for {module.__class__.__name__}', force=False)
                 if init_on_meta:
                     self.skip_init_depth -= 1
-            wrapped._ds_has_wrapped = True
+            wrapper._ds_has_wrapped = True
             return wrapper
 
         def _enable_class(cls):