pt: add necessary jit.export (#3337)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
deepmodeling · Feb 25, 2024 · 261c802 · 261c802
1 parent 91049df
commit 261c802
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 0 deletions.
diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py
@@ -55,6 +55,7 @@ def forward(
             model_predict["updated_coord"] += coord
         return model_predict
 
+    @torch.jit.export
     def forward_lower(
         self,
         extended_coord,

diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py
@@ -58,6 +58,7 @@ def forward(
             model_predict["force"] = model_ret["dforce"]
         return model_predict
 
+    @torch.jit.export
     def forward_lower(
         self,
         extended_coord,

diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
@@ -57,6 +57,7 @@ def forward(
             model_predict["updated_coord"] += coord
         return model_predict
 
+    @torch.jit.export
     def forward_lower(
         self,
         extended_coord,

diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
@@ -202,6 +202,7 @@ def forward_common_lower(
             )
             return model_predict
 
+        @torch.jit.export
         def format_nlist(
             self,
             extended_coord: torch.Tensor,

diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py
@@ -47,6 +47,7 @@ def forward(
             model_predict["updated_coord"] += coord
         return model_predict
 
+    @torch.jit.export
     def forward_lower(
         self,
         extended_coord,