fix

AlibabaPAI · Sep 22, 2024 · 3b948c5 · 3b948c5
1 parent effe7c1
commit 3b948c5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchacc/utils/optim_utils.py b/torchacc/utils/optim_utils.py
@@ -50,8 +50,8 @@ def get_layer_full_info(shard_metadata, model_state_dict):
             if name != 'model':
                 break
             else:
-                n = n + 1
-        name_splits = name_splits[n:]
+                model_num = model_num + 1
+        name_splits = name_splits[model_num:]
         name = ".".join(name_splits)
 
         for idx, sep in enumerate(name_splits):