From 835034ef3250c42ac4d21b0d3753a57571d3356d Mon Sep 17 00:00:00 2001 From: Fangjun Kuang Date: Fri, 6 May 2022 17:46:24 +0800 Subject: [PATCH] Save batch to disk on exception. --- 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 7a4b03cce..51c1a231a 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