[CI/Build] Simplify model loading for HfRunner (#5251)

2024-06-05 01:09:19 +08:00
parent 27208be66e
commit 9ba093b4f4
3 changed files with 18 additions and 13 deletions
--- a/tests/models/test_embedding.py
+++ b/tests/models/test_embedding.py
@ -28,7 +28,7 @@ def test_models(
    model: str,
    dtype: str,
 ) -> None:
-    hf_model = hf_runner(model, dtype=dtype)
+    hf_model = hf_runner(model, dtype=dtype, is_embedding_model=True)
    hf_outputs = hf_model.encode(example_prompts)
    del hf_model