Skip to content

Commit

Permalink
fix(Q1): merge #3836 in 2024Q1 (#3852)
Browse files Browse the repository at this point in the history
  • Loading branch information
iProzd authored Jun 4, 2024
1 parent 22026a8 commit d39bb94
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion deepmd/pt/utils/auto_batch_size.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,17 @@ def is_oom_error(self, e: Exception) -> bool:
e : Exception
Exception
"""
return isinstance(e, RuntimeError) and "CUDA out of memory." in e.args[0]
# several sources think CUSOLVER_STATUS_INTERNAL_ERROR is another out-of-memory error,
# such as https://github.com/JuliaGPU/CUDA.jl/issues/1924
# (the meaningless error message should be considered as a bug in cusolver)
if isinstance(e, RuntimeError) and (
"CUDA out of memory." in e.args[0]
or "cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR" in e.args[0]
):
# Release all unoccupied cached memory
torch.cuda.empty_cache()
return True
return False

def execute_all(
self, callable: Callable, total_size: int, natoms: int, *args, **kwargs
Expand Down

0 comments on commit d39bb94

Please sign in to comment.