fix(pt): improve out-of-memory capture

I just received another error message that reports out of memory. It's a bad design of PyTorch that all errors use a general `RuntimeError`.

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/pt/utils/auto_batch_size.py

@@ -57,6 +57,7 @@ def is_oom_error(self, e: Exception) -> bool:
         # (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 "CUDA driver error: out of memory" in e.args[0]
             or "cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR" in e.args[0]
         ):
             # Release all unoccupied cached memory