Print the reminder for the illegal memory error in the AutoBatchSize …

…under tf (#4283)

#3822 added a reminder for the illegal memory error. However, this
reminder is only needed for tf. This PR moves the illegal memory
reminder from base class AutoBatchSize to the inherited class under tf.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced `AutoBatchSize` class to initialize batch size from an
environment variable, improving user guidance on memory management with
TensorFlow.
- **Bug Fixes**
- Removed redundant logging during initialization to streamline the
process when GPU resources are available.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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.
 

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