From 9fbf7fc01c4683f104566b1f5911a7ffa4cca760 Mon Sep 17 00:00:00 2001
From: zhulinJulia24 <145004780+zhulinJulia24@users.noreply.github.com>
Date: Wed, 29 May 2024 11:00:49 +0800
Subject: [PATCH] Update test_hf_model.py

---
 tests/test_hf_model.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/test_hf_model.py b/tests/test_hf_model.py
index b6118a5..33162f4 100644
--- a/tests/test_hf_model.py
+++ b/tests/test_hf_model.py
@@ -127,7 +127,9 @@ class TestMath:
 
     @pytest.mark.parametrize(
         'model_name',
-        ['internlm/internlm2-math-7b', 'internlm/internlm2-math-base-7b'],
+        ['internlm/internlm2-math-7b', 'internlm/internlm2-math-base-7b',
+         'internlm/internlm2-math-plus-1_8b', 'internlm/internlm2-math-plus-7b'
+        ],
     )
     @pytest.mark.parametrize(
         'usefast',
@@ -163,7 +165,8 @@ class TestMMModel:
         'model_name',
         [
             'internlm/internlm-xcomposer2-7b',
-            'internlm/internlm-xcomposer2-7b-4bit'
+            'internlm/internlm-xcomposer2-7b-4bit',
+            'internlm/internlm-xcomposer2-4khd-7b'
         ],
     )
     def test_demo_default(self, model_name):