Update auto_batch_size.py

iProzd · Jan 7, 2025 · 5702a47 · 5702a47
1 parent 3feb0f6
commit 5702a47
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/deepmd/pt/utils/auto_batch_size.py b/deepmd/pt/utils/auto_batch_size.py
@@ -53,6 +53,7 @@ def is_oom_error(self, e: Exception) -> bool:
             "CUDA out of memory." in e.args[0]
             or "CUDA driver error: out of memory" in e.args[0]
             or "cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR" in e.args[0]
+            or "CUDA error: CUBLAS_STATUS_INTERNAL_ERROR" in e.args[0]
         ):
             # Release all unoccupied cached memory
             torch.cuda.empty_cache()