Update auto_batch_size.py

iProzd · iProzd · commit ba4e2ad7ddca · 2025-09-18T11:26:03.000+08:00
diff --git a/deepmd/pt/utils/auto_batch_size.py b/deepmd/pt/utils/auto_batch_size.py
@@ -57,6 +57,7 @@ def is_oom_error(self, e: Exception) -> bool:
                 or "cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR" in e.args[0]
                 # https://github.com/deepmodeling/deepmd-kit/issues/4594
                 or "CUDA error: out of memory" in e.args[0]
+                or "nonzero is not supported" in e.args[0]
             )
         ) or isinstance(e, torch.cuda.OutOfMemoryError):
             # Release all unoccupied cached memory

Original file line number	Diff line number	Diff line change
`@@ -57,6 +57,7 @@ def is_oom_error(self, e: Exception) -> bool:`
`57`	`57`	`or "cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR" in e.args[0]`
`58`	`58`	`# https://github.com/deepmodeling/deepmd-kit/issues/4594`
`59`	`59`	`or "CUDA error: out of memory" in e.args[0]`
	`60`	`+ or "nonzero is not supported" in e.args[0]`
`60`	`61`	`)`
`61`	`62`	`) or isinstance(e, torch.cuda.OutOfMemoryError):`
`62`	`63`	`# Release all unoccupied cached memory`