fix padding type casting

Blaizzy · Nov 26, 2024 · 0b002a9 · 0b002a9
1 parent 88c31b4
commit 0b002a9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mlx_vlm/utils.py b/mlx_vlm/utils.py
@@ -631,8 +631,8 @@ def _pad_vision_model(model: nn.Module, vision_size: int, divisor: int = 64) ->
                 )
 
                 if new_out != out_features or new_in != in_features:
-                    new_weight = mx.zeros((new_out, new_in))
-                    new_bias = mx.zeros((new_out))
+                    new_weight = mx.zeros((new_out, new_in), dtype=module.weight.dtype)
+                    new_bias = mx.zeros((new_out), dtype=module.bias.dtype)
 
                     new_weight[:out_features, :in_features] = module.weight
                     module.weight = new_weight