diff --git a/timm_image.py b/timm_image.py
index f74eb9a..bbd441f 100644
--- a/timm_image.py
+++ b/timm_image.py
@@ -100,7 +100,7 @@ class TimmImage(NNOperator):
             pass
 
     @torch_no_grad
-    def __call__(self, data: Union[List[towhee._types.Image], towhee._types.Image]):
+    def __call__(self, data: Union[List['towhee.types.Image'], 'towhee.types.Image']):
         if not isinstance(data, list):
             imgs = [data]
         else: