fix bug

wanghan-iapcm · Feb 11, 2024 · 69425af · 69425af
1 parent cb7b96d
commit 69425af
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/deepmd/pt/model/model/ener.py b/deepmd/pt/model/model/ener.py
@@ -66,8 +66,8 @@ def forward_lower(
         extended_atype,
         nlist,
         mapping: Optional[torch.Tensor] = None,
-        fparam: Optional[np.ndarray] = None,
-        aparam: Optional[np.ndarray] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
     ):
         model_ret = self.forward_common_lower(
@@ -148,6 +148,8 @@ def forward_lower(
         extended_atype,
         nlist,
         mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
     ):
         model_ret = self.forward_common_lower(