update

Author: dsikka

src/llmcompressor/modeling/moe_context.py

@@ -45,7 +45,7 @@ class MoECalibrationModule(ABC, torch.nn.Module):
 
     is_permanent: bool = False
 
-    def restore(self) -> torch.nn.Module:
+    def restore(self, original: torch.nn.Module) -> torch.nn.Module:
         """
         Restore the original module structure.
 
@@ -163,5 +163,5 @@ def moe_calibration_context(
         # Step 2: Restore non-permanent modules
         for name, (original, replacement) in replaced.items():
             if not replacement.is_permanent:
-                restored = replacement.restore()
+                restored = replacement.restore(original)
                 model.set_submodule(name, restored)

src/llmcompressor/modeling/qwen3_moe.py

@@ -98,6 +98,9 @@ def forward(self, hidden_states: torch.Tensor):
         )
         return final_hidden_states, router_logits
 
+    def restore(self, original: torch.nn.Module) -> torch.nn.Module:
+        return original
+
 
 # Legacy function for backward compatibility
 def replace(

src/llmcompressor/modeling/qwen3_next_moe.py

@@ -115,6 +115,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         )
         return final_hidden_states, router_logits
 
+    def restore(self, original: torch.nn.Module) -> torch.nn.Module:
+        return original
+
 
 def replace(
     config,