Correct name of MixtralBlockSparseTop2MLP (L -> l) (#1667)

axolotl-ai-cloud · May 28, 2024 · 65db903 · 65db903
1 parent 6a5a725
commit 65db903
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/axolotl/monkeypatch/mixtral/__init__.py b/src/axolotl/monkeypatch/mixtral/__init__.py
@@ -42,9 +42,9 @@ def moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return final_hidden_states, router_logits
 
     from transformers.models.mixtral.modeling_mixtral import (
-        MixtralBLockSparseTop2MLP,
+        MixtralBlockSparseTop2MLP,
         MixtralSparseMoeBlock,
     )
 
-    MixtralBLockSparseTop2MLP.forward = mlp_forward
+    MixtralBlockSparseTop2MLP.forward = mlp_forward
     MixtralSparseMoeBlock.forward = moe_forward