mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-08-08 17:42:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Save batch to disk on exception. (#350)

Browse Source
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
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
egs/librispeech/ASR/pruned_transducer_stateless2/train.py
View File

@ -1001,7 +1001,7 @@ def scan_pessimistic_batches_for_oom(
loss.backward()
optimizer.step()
optimizer.zero_grad()
except RuntimeError as e:
except Exception as e:
if "CUDA out of memory" in str(e):
logging.error(
"Your GPU ran out of memory with the current "
@ -1010,7 +1010,7 @@ def scan_pessimistic_batches_for_oom(
f"Failing criterion: {criterion} "
f"(={crit_values[criterion]}) ..."
)
display_and_save_batch(batch, params=params, sp=sp)
display_and_save_batch(batch, params=params, sp=sp)
raise