diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py
index c2fd81a02e..a300b8a69d 100644
--- a/deepmd/pd/model/descriptor/se_a.py
+++ b/deepmd/pd/model/descriptor/se_a.py
@@ -19,6 +19,7 @@
     prod_env_mat,
 )
 from deepmd.pd.utils import (
+    decomp,
     env,
 )
 from deepmd.pd.utils.env import (
@@ -744,7 +745,8 @@ def forward(
                     "Compressed environment is not implemented yet."
                 )
             else:
-                if rr.numel() > 0:
+                # NOTE: control flow with double backward is not supported well yet by paddle.jit
+                if not paddle.framework.in_dynamic_mode() or decomp.numel(rr) > 0:
                     rr = rr * mm.unsqueeze(2).astype(rr.dtype)
                     ss = rr[:, :, :1]
                     if self.compress:
diff --git a/deepmd/pd/model/network/layernorm.py b/deepmd/pd/model/network/layernorm.py
index 4d37b208f9..22f9d16c13 100644
--- a/deepmd/pd/model/network/layernorm.py
+++ b/deepmd/pd/model/network/layernorm.py
@@ -100,8 +100,8 @@ def forward(
         yy: paddle.Tensor
             The output.
         """
-        # if xx.numel() > 0:
-        if decomp.numel(xx):
+        # NOTE: control flow with double backward is not supported well yet by paddle.jit
+        if not paddle.framework.in_dynamic_mode() or decomp.numel(xx) > 0:
             variance, mean = (
                 paddle.var(xx, axis=-1, unbiased=False, keepdim=True),
                 paddle.mean(xx, axis=-1, keepdim=True),
diff --git a/deepmd/pd/utils/nlist.py b/deepmd/pd/utils/nlist.py
index ae9db628a1..3f6f99bc12 100644
--- a/deepmd/pd/utils/nlist.py
+++ b/deepmd/pd/utils/nlist.py
@@ -7,6 +7,7 @@
 import paddle
 
 from deepmd.pd.utils import (
+    decomp,
     env,
 )
 from deepmd.pd.utils.region import (
@@ -97,7 +98,9 @@ def build_neighbor_list(
     nall = coord.shape[1] // 3
     # fill virtual atoms with large coords so they are not neighbors of any
     # real atom.
-    if coord.numel() > 0:
+
+    # NOTE: control flow with double backward is not supported well yet by paddle.jit
+    if not paddle.framework.in_dynamic_mode() or decomp.numel(coord) > 0:
         xmax = paddle.max(coord) + 2.0 * rcut
     else:
         xmax = paddle.zeros([], dtype=coord.dtype).to(device=coord.place) + 2.0 * rcut
@@ -240,7 +243,8 @@ def build_directional_neighbor_list(
     nall_neig = coord_neig.shape[1] // 3
     # fill virtual atoms with large coords so they are not neighbors of any
     # real atom.
-    if coord_neig.numel() > 0:
+    # NOTE: control flow with double backward is not supported well yet by paddle.jit
+    if not paddle.framework.in_dynamic_mode() or decomp.numel(coord_neig) > 0:
         xmax = paddle.max(coord_cntl) + 2.0 * rcut
     else:
         xmax = (