diff --git a/deepmd/infer/model_devi.py b/deepmd/infer/model_devi.py
index c214e39e92..a37dfd34c5 100644
--- a/deepmd/infer/model_devi.py
+++ b/deepmd/infer/model_devi.py
@@ -303,7 +303,7 @@ def calc_model_devi(
     energies = []
     forces = []
     virials = []
-    natom = atype.shape[-1]
+    natom = np.array(atype).shape[-1]
     for dp in models:
         ret = dp.eval(
             coord,