fix gemma embedding scaling

Blaizzy · May 24, 2024 · c70c916 · c70c916
1 parent c25d9c0
commit c70c916
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mlx_vlm/models/paligemma/language.py b/mlx_vlm/models/paligemma/language.py
@@ -152,10 +152,12 @@ def __call__(
         # for passing merged input embeddings
         if inputs_embeds is None:
             h = self.embed_tokens(inputs)
-            h = h * (self.args.hidden_size**0.5)
+
         else:
             h = inputs_embeds
 
+        h = h * (self.args.hidden_size**0.5)
+
         if cache is not None:
             mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
             mask = mask.astype(h.dtype)