From 5702a47ac2287b02b176a46d0f7914b9f36a386b Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 7 Jan 2025 15:03:14 +0800 Subject: [PATCH] Update auto_batch_size.py --- deepmd/pt/utils/auto_batch_size.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/pt/utils/auto_batch_size.py b/deepmd/pt/utils/auto_batch_size.py index 53942a176a..f1351a1fd2 100644 --- 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()