From a0dbfba77dd64bfb4c2980ffa2283a6eb1050e79 Mon Sep 17 00:00:00 2001 From: Fangjun Kuang Date: Thu, 5 May 2022 12:47:24 +0800 Subject: [PATCH] Fixes after review. --- egs/librispeech/ASR/pruned_transducer_stateless2/train.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/train.py b/egs/librispeech/ASR/pruned_transducer_stateless2/train.py index bde297ae4..58604ab43 100755 --- a/egs/librispeech/ASR/pruned_transducer_stateless2/train.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless2/train.py @@ -691,9 +691,8 @@ def train_one_epoch( scaler.step(optimizer) scaler.update() optimizer.zero_grad() - except RuntimeError as e: - if "CUDA out of memory" in str(e): - display_and_save_batch(batch, params=params, sp=sp) + except: + display_and_save_batch(batch, params=params, sp=sp) raise if params.print_diagnostics and batch_idx == 5: