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()