diff --git a/vit_image_embedding.py b/vit_image_embedding.py
index 8812503..c48566e 100644
--- a/vit_image_embedding.py
+++ b/vit_image_embedding.py
@@ -19,10 +19,11 @@ from pathlib import Path
 from PIL import Image
 import torch
 import numpy
-
+import os
 
 from towhee.operator import Operator
-
+from timm.data import resolve_data_config
+from timm.data.transforms_factory import create_transform
 
 class VitImageEmbedding(Operator):
     """
@@ -37,13 +38,16 @@ class VitImageEmbedding(Operator):
     def __init__(self, model_name: str = 'vit_large_patch16_224',
                  framework: str = 'pytorch', weights_path: str = None) -> None:
         super().__init__()
-        sys.path.append(str(Path(__file__).parent))
         if framework == 'pytorch':
-            import pytorch
-            from pytorch.model import Model
-        self.model = Model(model_name, weights_path)
-        config = pytorch.resolve_data_config({}, model=self.model._model)
-        self.tfms = pytorch.create_transform(**config)
+            import importlib.util
+            path = os.path.join(str(Path(__file__).parent), 'pytorch', 'model.py')
+            opname = os.path.basename(str(Path(__file__))).split('.')[0]
+            spec = importlib.util.spec_from_file_location(opname, path)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+        self.model = module.Model(model_name, weights_path)
+        config = resolve_data_config({}, model=self.model._model)
+        self.tfms = create_transform(**config)
 
     def __call__(self, img_path: str) -> NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]):
         Outputs = NamedTuple('Outputs', [('feature_vector', numpy.ndarray)])