diff --git a/deepmd/env.py b/deepmd/env.py
index b6dd3e241f..95f29c244b 100644
--- a/deepmd/env.py
+++ b/deepmd/env.py
@@ -274,6 +274,8 @@ def get_tf_session_config() -> Any:
     """
     set_tf_default_nthreads()
     intra, inter = get_tf_default_nthreads()
+    if int(os.environ.get("DP_JIT", 0)):
+        set_env_if_empty("TF_XLA_FLAGS", "--tf_xla_auto_jit=2")
     config = tf.ConfigProto(
         gpu_options=tf.GPUOptions(allow_growth=True),
         intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter
diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md
index 0fef1fa6a3..7c7247a36b 100644
--- a/doc/train/training-advanced.md
+++ b/doc/train/training-advanced.md
@@ -137,6 +137,7 @@ One can set other environmental variables:
 | --------------------- | ---------------------- | ------------- | -------------------------- |
 | DP_INTERFACE_PREC     | `high`, `low`          | `high`        | Control high (double) or low (float) precision of training. |
 | DP_AUTO_PARALLELIZATION | 0, 1                 | 0             | Enable auto parallelization for CPU operators. |
+| DP_JIT                | 0, 1                   | 0             | Enable JIT. Note that this option may either improve or decrease the performance. Requires TensorFlow supports JIT.  |
 
 
 ## Adjust `sel` of a frozen model