only show the reminder when GPU is used

deepmodeling · Oct 30, 2024 · 5870a93 · 5870a93
1 parent 94b2d41
commit 5870a93
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/deepmd/tf/utils/batch_size.py b/deepmd/tf/utils/batch_size.py
@@ -23,11 +23,12 @@ 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:
-            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). "
-            )
+            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.