From 403d1744fff4834bf316b1fd715bdceda62646fa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20=C5=BBelasko?= <petezor@gmail.com>
Date: Fri, 15 Oct 2021 10:05:13 -0400
Subject: [PATCH] Introduce backprop in finding OOM batches

---
 egs/librispeech/ASR/conformer_ctc/train.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/egs/librispeech/ASR/conformer_ctc/train.py b/egs/librispeech/ASR/conformer_ctc/train.py
index 9d3de020e..db167fd9d 100755
--- a/egs/librispeech/ASR/conformer_ctc/train.py
+++ b/egs/librispeech/ASR/conformer_ctc/train.py
@@ -618,13 +618,17 @@ def run(rank, world_size, args):
         )
         batch = train_dl.dataset[cuts]
         try:
-            compute_loss(
+            optimizer.zero_grad()
+            loss, _ = compute_loss(
                 params=params,
                 model=model,
                 batch=batch,
                 graph_compiler=graph_compiler,
                 is_training=True,
             )
+            loss.backward()
+            clip_grad_norm_(model.parameters(), 5.0, 2.0)
+            optimizer.step()
             logging.info("OK!")
         except RuntimeError as e:
             if "CUDA out of memory" in str(e):