fix type hint

deepmodeling · Jan 31, 2024 · 96bbb45 · 96bbb45
1 parent 0c1913b
commit 96bbb45
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/deepmd/model_format/atomic_model.py b/deepmd/model_format/atomic_model.py
@@ -10,8 +10,6 @@
     Optional,
 )
 
-import numpy as np
-
 from .output_def import (
     FittingOutputDef,
 )
@@ -44,8 +42,8 @@ def forward_atomic(
             extended_atype: T_Tensor,
             nlist: T_Tensor,
             mapping: Optional[T_Tensor] = None,
-            fparam: Optional[np.ndarray] = None,
-            aparam: Optional[np.ndarray] = None,
+            fparam: Optional[T_Tensor] = None,
+            aparam: Optional[T_Tensor] = None,
         ) -> Dict[str, T_Tensor]:
             raise NotImplementedError