Save batch to disk on exception. (#350)

k2-fsa · May 6, 2022 · 32f05c0 · 32f05c0
1 parent 00c48ec
commit 32f05c0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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