diff --git a/tests/unit/encoders/test_vit.py b/tests/unit/encoders/test_vit.py index 72d01b349bc83a0b824e2f2f746d29c0553c94ce..9970e31075a44396c5fbfeaff0a2ce0a846487ca 100644 --- a/tests/unit/encoders/test_vit.py +++ b/tests/unit/encoders/test_vit.py @@ -4,8 +4,14 @@ import torch from PIL import Image from semantic_router.encoders import VitEncoder +from transformers import AutoTokenizer, AutoModel test_model_name = "hf-internal-testing/tiny-random-vit" + +# force the model download +tokenizer = AutoTokenizer.from_pretrained(test_model_name, force_download=True) +model = AutoModel.from_pretrained(test_model_name, force_download=True) + vit_encoder = VitEncoder(name=test_model_name) embed_dim = 32