diff --git a/src/sparseml/modifiers/utils/layer_compressor.py b/src/sparseml/modifiers/utils/layer_compressor.py
index 3818d366736..a614789ff8c 100644
--- 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