From 21f49bab444e0761ab7fc7ed20b3c81fb7924d17 Mon Sep 17 00:00:00 2001
From: Jake Zhang <ake020675@163.com>
Date: Tue, 16 Apr 2024 11:55:10 +0800
Subject: [PATCH] Fix: Importing ViT in `gen_model.py` (#80)

---
 python/uform/gen_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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)