delete torch.jit.export of get_default_fparam (#51)

Author: Chengqian-Zhang

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -5,7 +5,6 @@
     Optional,
 )
 
-import numpy as np
 import torch
 
 from deepmd.dpmodel import (
@@ -366,7 +365,7 @@ def get_dim_fparam(self) -> int:
     def has_default_fparam(self) -> bool:
         return self.fitting_net.has_default_fparam()
 
-    def get_default_fparam(self) -> Optional[np.array]:
+    def get_default_fparam(self) -> Optional[torch.Tensor]:
         return self.fitting_net.get_default_fparam()
 
     def get_dim_aparam(self) -> int:

deepmd/pt/model/model/make_model.py

@@ -3,7 +3,6 @@
     Optional,
 )
 
-import numpy as np
 import torch
 
 from deepmd.dpmodel import (
@@ -531,8 +530,7 @@ def get_dim_fparam(self) -> int:
         def has_default_fparam(self) -> bool:
             return self.atomic_model.has_default_fparam()
 
-        @torch.jit.export
-        def get_default_fparam(self) -> Optional[np.array]:
+        def get_default_fparam(self) -> Optional[torch.Tensor]:
             return self.atomic_model.get_default_fparam()
 
         @torch.jit.export

deepmd/pt/model/task/fitting.py

@@ -619,8 +619,8 @@ def get_dim_fparam(self) -> int:
     def has_default_fparam(self) -> bool:
         return self.default_fparam is not None
 
-    def get_default_fparam(self) -> Optional[np.array]:
-        return self.default_fparam_tensor.cpu().numpy()
+    def get_default_fparam(self) -> Optional[torch.Tensor]:
+        return self.default_fparam_tensor
 
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""

deepmd/pt/train/training.py

@@ -1305,7 +1305,9 @@ def get_additional_data_requirement(_model):
     additional_data_requirement = []
     if _model.get_dim_fparam() > 0:
         _fparam_default = (
-            _model.get_default_fparam() if _model.has_default_fparam() else 0.0
+            _model.get_default_fparam().cpu().numpy()
+            if _model.has_default_fparam()
+            else 0.0
         )
         fparam_requirement_items = [
             DataRequirementItem(