Fix: Importing ViT in gen_model.py (#80)

unum-cloud · Apr 16, 2024 · 21f49ba · 21f49ba
1 parent a5d84fc
commit 21f49ba
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/uform/gen_model.py b/python/uform/gen_model.py
@@ -15,7 +15,7 @@
 from transformers.processing_utils import ProcessorMixin
 from transformers.tokenization_utils_base import BatchEncoding
 
-from uform.models import VisualEncoder
+from uform.torch_models import VisualEncoder
 
 IMAGENET_MEAN = (0.48145466, 0.4578275, 0.40821073)
 IMAGENET_STD = (0.26862954, 0.26130258, 0.27577711)