From 613c5bb9458f9ced270fd98c9d877e00c0ad165b Mon Sep 17 00:00:00 2001
From: Cyrus Leung <tlleungac@connect.ust.hk>
Date: Sat, 15 Mar 2025 00:11:23 +0800
Subject: [PATCH] [Bugfix] Fix Aria test loading (#14823)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
---
 tests/models/decoder_only/vision_language/test_models.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/models/decoder_only/vision_language/test_models.py b/tests/models/decoder_only/vision_language/test_models.py
index 84a5260ad9a08..a0f1229f0af5a 100644
--- a/tests/models/decoder_only/vision_language/test_models.py
+++ b/tests/models/decoder_only/vision_language/test_models.py
@@ -9,7 +9,8 @@ from pathlib import PosixPath
 
 import pytest
 from packaging.version import Version
-from transformers import AutoModelForPreTraining, AutoModelForVision2Seq
+from transformers import (AutoModelForImageTextToText, AutoModelForPreTraining,
+                          AutoModelForVision2Seq)
 from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.platforms import current_platform
@@ -163,6 +164,7 @@ VLM_TEST_SETTINGS = {
         img_idx_to_prompt=lambda idx: "<fim_prefix><|img|><fim_suffix>\n",
         max_model_len=4096,
         max_num_seqs=2,
+        auto_cls=AutoModelForImageTextToText,
         single_image_prompts=IMAGE_ASSETS.prompts({
             "stop_sign": "<vlm_image>Please describe the image shortly.",
             "cherry_blossom": "<vlm_image>Please infer the season with reason.",