From 96bbb454f338570d1989af3519205297e1cab72b Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Wed, 31 Jan 2024 22:37:48 +0800
Subject: [PATCH] fix type hint

---
 deepmd/model_format/atomic_model.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/deepmd/model_format/atomic_model.py b/deepmd/model_format/atomic_model.py
index 5ec2c1a67c..cdc08fe1f6 100644
--- 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