From 32f05c00e3162a6e5d06fb620b9918a23de8035c Mon Sep 17 00:00:00 2001 From: Fangjun Kuang Date: Fri, 6 May 2022 17:49:40 +0800 Subject: [PATCH] Save batch to disk on exception. (#350) --- egs/librispeech/ASR/pruned_transducer_stateless2/train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/train.py b/egs/librispeech/ASR/pruned_transducer_stateless2/train.py index 7a4b03cce7..51c1a231a4 100755 --- a/egs/librispeech/ASR/pruned_transducer_stateless2/train.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless2/train.py @@ -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