diff --git a/python/uform/gen_model.py b/python/uform/gen_model.py
index 35faae1..c03b6eb 100644
--- 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)