fix jit

Author: iProzd

deepmd/dpmodel/utils/network.py

@@ -646,8 +646,10 @@ def call_until_last(self, x):
             np.ndarray
                 The output before last layer.
             """
-            for layer in self.layers[:-1]:
-                x = layer(x)
+            # avoid slice (self.layers[:-1]) for jit
+            for ii, layer in enumerate(self.layers):
+                if ii < len(self.layers) - 1:
+                    x = layer(x)
             return x
 
         def clear(self) -> None:

deepmd/pt/infer/deep_eval.py

@@ -131,8 +131,8 @@ def __init__(
                 state_dict = state_dict_head
             model = get_model(self.input_param).to(DEVICE)
             # TODO fix jit
-            # if not self.input_param.get("hessian_mode"):
-            #     model = torch.jit.script(model)
+            if not self.input_param.get("hessian_mode"):
+                model = torch.jit.script(model)
             self.dp = ModelWrapper(model)
             self.dp.load_state_dict(state_dict)
         elif str(self.model_path).endswith(".pth"):

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -271,7 +271,7 @@ def forward_atomic(
         )
         if self.enable_eval_fitting_last_layer_hook:
             assert "middle_output" in fit_ret, (
-                f"eval_fitting_last_layer not supported for fitting net {type(self.fitting_net.__class__)}!"
+                "eval_fitting_last_layer not supported for this fitting net!"
             )
             self.eval_fitting_last_layer_list.append(
                 fit_ret.pop("middle_output").detach()