[Bugfix] Fix FusedMoE LoRA w2_output_size (#31949)

Signed-off-by: Xin Yang <xyangx@amazon.com>
2026-01-09 00:54:05 -05:00
parent 29ce48221c
commit 707b240d7e
1 changed files with 1 additions and 1 deletions
@@ -720,7 +720,7 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
        """
        Full size
        """
-        return self.w2_lora_a_stacked[0].shape[-2]
+        return self.base_layer.hidden_size

    @classmethod
    def can_replace_layer(