only gather the current transformer layer (#2173) (#2174)

neuralmagic · Mar 12, 2024 · 7d2f7ce · 7d2f7ce
1 parent 9ae35bd
commit 7d2f7ce
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/sparseml/modifiers/utils/layer_compressor.py b/src/sparseml/modifiers/utils/layer_compressor.py
@@ -122,8 +122,7 @@ def revert_layer_wrappers(self):
         Reverts wrapped root modules back to their original structure
         """
         for name, module_wrapper in self.modules.items():
-            full_name = self._get_full_submodule_name(name)
-            set_layer(full_name, module_wrapper.layer, self.model)
+            set_layer(name, module_wrapper.layer, self.layer)
             module_wrapper.free()
         self.modules = None