fix encoder states

Blaizzy · May 22, 2024 · c25d9c0 · c25d9c0
1 parent d88d1af
commit c25d9c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mlx_vlm/models/paligemma/vision.py b/mlx_vlm/models/paligemma/vision.py
@@ -162,7 +162,7 @@ def __call__(
         for l in self.layers:
             x = l(x, mask=mask)
             if output_hidden_states:
-                encoder_states = encoder_states + (h,)
+                encoder_states = encoder_states + (x,)
 
             h = x[0]