fix mamba models conversion (#1065)

ml-explore · Oct 22, 2024 · 9000e28 · 9000e28
1 parent d1d4808
commit 9000e28
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/llms/mlx_lm/models/mamba.py b/llms/mlx_lm/models/mamba.py
@@ -205,7 +205,7 @@ def __call__(self, inputs: mx.array, cache=None):
 
     def sanitize(self, weights):
         for k, v in weights.items():
-            if "conv1d.weight" in k and v.ndim == 3:
+            if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)
         return weights
 

diff --git a/llms/mlx_lm/models/recurrent_gemma.py b/llms/mlx_lm/models/recurrent_gemma.py
@@ -440,7 +440,7 @@ def layers(self):
 
     def sanitize(self, weights):
         for k, v in weights.items():
-            if "conv_1d.weight" in k and v.ndim == 3:
+            if "conv_1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)
         if "lm_head.weight" not in weights:
             self.pop("lm_head")