diff --git a/deepmd/tf/utils/batch_size.py b/deepmd/tf/utils/batch_size.py
index 33f1ec0da0..438bf36703 100644
--- a/deepmd/tf/utils/batch_size.py
+++ b/deepmd/tf/utils/batch_size.py
@@ -1,4 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import os
+
 from packaging.version import (
     Version,
 )
@@ -11,9 +13,23 @@
     OutOfMemoryError,
 )
 from deepmd.utils.batch_size import AutoBatchSize as AutoBatchSizeBase
+from deepmd.utils.batch_size import (
+    log,
+)
 
 
 class AutoBatchSize(AutoBatchSizeBase):
+    def __init__(self, initial_batch_size: int = 1024, factor: float = 2.0) -> None:
+        super().__init__(initial_batch_size, factor)
+        DP_INFER_BATCH_SIZE = int(os.environ.get("DP_INFER_BATCH_SIZE", 0))
+        if not DP_INFER_BATCH_SIZE > 0:
+            if self.is_gpu_available():
+                log.info(
+                    "If you encounter the error 'an illegal memory access was encountered', this may be due to a TensorFlow issue. "
+                    "To avoid this, set the environment variable DP_INFER_BATCH_SIZE to a smaller value than the last adjusted batch size. "
+                    "The environment variable DP_INFER_BATCH_SIZE controls the inference batch size (nframes * natoms). "
+                )
+
     def is_gpu_available(self) -> bool:
         """Check if GPU is available.
 
diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py
index 259fe93bdb..5ab06e55e2 100644
--- a/deepmd/utils/batch_size.py
+++ b/deepmd/utils/batch_size.py
@@ -61,11 +61,6 @@ def __init__(self, initial_batch_size: int = 1024, factor: float = 2.0) -> None:
             self.maximum_working_batch_size = initial_batch_size
             if self.is_gpu_available():
                 self.minimal_not_working_batch_size = 2**31
-                log.info(
-                    "If you encounter the error 'an illegal memory access was encountered', this may be due to a TensorFlow issue. "
-                    "To avoid this, set the environment variable DP_INFER_BATCH_SIZE to a smaller value than the last adjusted batch size. "
-                    "The environment variable DP_INFER_BATCH_SIZE controls the inference batch size (nframes * natoms). "
-                )
             else:
                 self.minimal_not_working_batch_size = (
                     self.maximum_working_batch_size + 1