Save batch to disk on exception. (#350)

This commit is contained in:
Fangjun Kuang 2022-05-06 17:49:40 +08:00 committed by GitHub
parent 00c48ec1f3
commit 32f05c00e3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1001,7 +1001,7 @@ def scan_pessimistic_batches_for_oom(
loss.backward() loss.backward()
optimizer.step() optimizer.step()
optimizer.zero_grad() optimizer.zero_grad()
except RuntimeError as e: except Exception as e:
if "CUDA out of memory" in str(e): if "CUDA out of memory" in str(e):
logging.error( logging.error(
"Your GPU ran out of memory with the current " "Your GPU ran out of memory with the current "