From 80ea2c05c27a5fada0d034951826d44d4738b9e9 Mon Sep 17 00:00:00 2001
From: Joao Gante <joaofranciscocardosogante@gmail.com>
Date: Wed, 23 Apr 2025 17:54:12 +0100
Subject: [PATCH] [tests, `qwen2_5_omni`] fix flaky tests (#37721)

---
 tests/generation/test_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
index 2bc1c351b8..1aa88abcb9 100644
--- a/tests/generation/test_utils.py
+++ b/tests/generation/test_utils.py
@@ -134,7 +134,7 @@ VLM_CLASS_NAMES = [
     "mistral3",
     "chameleon",
     "internvl",
-    "qwen2_5_omni",
+    "qwen2_5omni",  # the file is named `qwen2_5_omni`, but the model class is `Qwen2_5Omni`
 ]
 
 
@@ -230,6 +230,7 @@ class GenerationTesterMixin:
             for key in [
                 "image_token_id",
                 "video_token_id",
+                "audio_token_id",
                 "vision_start_token_id",
                 "audio_start_token_id",
                 "audio_end_token_id",