Fix image tests. (#1253)

Signed-off-by: Qiliang Cui <derrhein@gmail.com>

Author: QiliangCui

tests/runner/test_multimodal_manager.py

@@ -144,8 +144,8 @@ def test_execute_mm_encoder_single_image(self):
         assert passed_pixel_values.dtype == jnp.bfloat16
 
         # Convert torch tensor for comparison
-        expected_pixel_values = dummy_pixel_values.unsqueeze(0).unsqueeze(
-            0).to(torch.float32).numpy().astype(jnp.bfloat16)
+        expected_pixel_values = dummy_pixel_values.unsqueeze(0).to(
+            torch.float32).numpy().astype(jnp.bfloat16)
         np.testing.assert_array_equal(np.asarray(passed_pixel_values),
                                       expected_pixel_values)
 
@@ -249,11 +249,10 @@ def test_execute_mm_encoder_multiple_images(self):
         assert "pixel_values" in kwargs_arg
 
         passed_pixel_values = kwargs_arg['pixel_values']
-        assert passed_pixel_values.shape == (2, 1, 3, 224, 224)
+        assert passed_pixel_values.shape == (2, 3, 224, 224)
 
-        expected_pixel_values = torch.stack(
-            [px_1, px_2],
-            dim=0).unsqueeze(1).to(torch.float32).numpy().astype(jnp.bfloat16)
+        expected_pixel_values = torch.stack([px_1, px_2], dim=0).to(
+            torch.float32).numpy().astype(jnp.bfloat16)
         np.testing.assert_array_equal(np.asarray(passed_pixel_values),
                                       expected_pixel_values)