Update make_hessian_model.py

Signed-off-by: Anchor Yu <91590308+1azyking@users.noreply.github.com>
deepmodeling · Dec 6, 2024 · d1361ca · d1361ca
1 parent c699e08
commit d1361ca
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/deepmd/pt/model/model/make_hessian_model.py b/deepmd/pt/model/model/make_hessian_model.py
@@ -115,8 +115,8 @@ def forward_common(
 
         def _cal_hessian_all(
             self,
-            coord,
-            atype,
+            coord: torch.Tensor,
+            atype: torch.Tensor,
             box: Optional[torch.Tensor] = None,
             fparam: Optional[torch.Tensor] = None,
             aparam: Optional[torch.Tensor] = None,
@@ -184,7 +184,7 @@ def __init__(
             self,
             obj: CM,
             ci: int,
-            atype,
+            atype: torch.Tensor,
             box: Optional[torch.Tensor],
             fparam: Optional[torch.Tensor],
             aparam: Optional[torch.Tensor],