fix

Author: HydrogenSulfate

deepmd/pd/entrypoints/main.py

@@ -368,9 +368,9 @@ def freeze(
         model.forward = paddle.jit.to_static(
             model.forward,
             input_spec=[
-                InputSpec([1, -1, 3], dtype="float64", name="coord"),  # coord
-                InputSpec([1, -1], dtype="int64", name="atype"),  # atype
-                InputSpec([1, 9], dtype="float64", name="box"),  # box
+                InputSpec([-1, -1, 3], dtype="float64", name="coord"),  # coord
+                InputSpec([-1, -1], dtype="int64", name="atype"),  # atype
+                InputSpec([-1, 9], dtype="float64", name="box"),  # box
                 None,  # fparam
                 None,  # aparam
                 # InputSpec([], dtype="bool", name="do_atomic_virial"),  # do_atomic_virial

deepmd/pd/model/descriptor/repflows.py

@@ -663,6 +663,8 @@ def forward(
                         place=paddle.CPUPlace(),
                     ),  # should be int of c++, placed on cpu
                 )
+                if not paddle.in_dynamic_mode():
+                    ret = paddle.assign(ret)
                 node_ebd_ext = ret.unsqueeze(0)
                 if has_spin:
                     node_ebd_real_ext, node_ebd_virtual_ext = paddle.split(

deepmd/pd/model/descriptor/repformers.py

@@ -548,6 +548,8 @@ def forward(
                         place=paddle.CPUPlace(),
                     ),  # should be int of c++, placed on cpu
                 )
+                if not paddle.in_dynamic_mode():
+                    ret = paddle.assign(ret)
                 g1_ext = ret.unsqueeze(0)
                 if has_spin:
                     g1_real_ext, g1_virtual_ext = paddle.split(

deepmd/pd/utils/serialization.py

@@ -69,9 +69,9 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
         model.forward,
         full_graph=True,
         input_spec=[
-            InputSpec([1, -1, 3], dtype="float64", name="coord"),
-            InputSpec([1, -1], dtype="int64", name="atype"),
-            InputSpec([1, 9], dtype="float64", name="box"),
+            InputSpec([-1, -1, 3], dtype="float64", name="coord"),
+            InputSpec([-1, -1], dtype="int64", name="atype"),
+            InputSpec([-1, 9], dtype="float64", name="box"),
             None,
             None,
             True,
@@ -88,9 +88,9 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
         model.forward_lower,
         full_graph=True,
         input_spec=[
-            InputSpec([1, -1, 3], dtype="float64", name="extended_coord"),
-            InputSpec([1, -1], dtype="int32", name="extended_atype"),
-            InputSpec([1, -1, -1], dtype="int32", name="nlist"),
+            InputSpec([-1, -1, 3], dtype="float64", name="extended_coord"),
+            InputSpec([-1, -1], dtype="int32", name="extended_atype"),
+            InputSpec([-1, -1, -1], dtype="int32", name="nlist"),
             None,
             None,
             None,
@@ -101,4 +101,5 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
     paddle.jit.save(
         model,
         model_file.split(".json")[0],
+        skip_prune_program=True,
     )