diff --git a/src/stannum/tube.py b/src/stannum/tube.py
index 905d796..3402e91 100644
--- a/src/stannum/tube.py
+++ b/src/stannum/tube.py
@@ -2,8 +2,7 @@
 import torch
 from taichi.lang.impl import axes
 from typing import Optional, Callable, Union, Tuple, List, Iterable, Dict, Any
-from .utils import to_taichi_type, need_auto_clearing_fields
-from taichi._lib.core.taichi_core import DataType as TiDataType
+from .utils import to_taichi_type, need_auto_clearing_fields, TiDataType
 from taichi.lang.field import ScalarField
 from taichi.lang.matrix import MatrixField
 from functools import partial
diff --git a/src/stannum/utils.py b/src/stannum/utils.py
index d10be02..adba30a 100644
--- a/src/stannum/utils.py
+++ b/src/stannum/utils.py
@@ -5,7 +5,11 @@
 from taichi.lang.kernel_impl import Kernel
 from taichi.lang.matrix import MatrixField
 from taichi.lang.field import ScalarField
-from taichi._lib.core.taichi_core import DataType as TiDataType
+
+if __ti_version >= (1, 1, 0):
+    from taichi._lib.core.taichi_python import DataType as TiDataType
+else:
+    from taichi._lib.core.taichi_core import DataType as TiDataType
 import torch
 import taichi as ti