From b566f3b9254bab734ca8117cfb48f1a95c8aad9f Mon Sep 17 00:00:00 2001
From: dohe0342 <kimdohe1070@gmail.com>
Date: Sat, 10 Dec 2022 14:01:43 +0900
Subject: [PATCH] from local

---
 .../.train.py.swp                             | Bin 98304 -> 98304 bytes
 .../train.py                                  |  20 +++++++++++++-----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.train.py.swp b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.train.py.swp
index 2afb5b7b4e1ddc96b2d00baf36b823b9825701b5..dc73f9efb89b8c8a326b446ba009e998cd49596b 100644
GIT binary patch
delta 590
zcmY+<Pe__!90%~<<KEXYvwV}39BuC<SOrap9ma@SbZEq&Lo^}O)Kmnuccy3<C<I-`
zI_!xa=pwYkVAMef91%2k@X`)uJMGl1uC^TPw7#!GV-I|Ie$T^?-=9LkSST3n0oUTI
z*Xr{z7vWfm9DO@M!)?j_Q(^AShEi@Q%;vQ;7*YH|JHss$jpg!5GtoHQbP<(c3zpyo
zC~##W+J!gZgkPOR7jOs%umcM)4KKk2zdMMUa1C|%1VM1aGq6L`NK`XYp1vWdz!AKM
zIhX<)T*^dWAP)hUfJg90BKiTJ;T>oYfC(6cCU*P|CHM&Ypzl2gPoW#kP&bs$XPM9z
zm?XZ2nZ(z#z2Y>@M)|R)Tb!hsl_#@y{YU0+vkpBl>p>a86Imt|htC)L`SynaepWT{
z<%pwoZ$$b;Z-xyBZ;rWmEi(By6pMz{#A-^_{7HX&Z7iOSrJ}x7EftNc{=|wH$+Mw<
z1ts4~DA=m%E?OjCDA9iW-@@9dXANTtUa8vnM#L)XnZu$d!<?-Ot@h$!hP7+lPs8`t
SyTo4)GfVfq%O^QDBi#Y(?UU~S

delta 307
zcmXBPF+_xM9LDkIulL@!>zwX&d6&qM*8HO-tx-4h5+^m(aMAp_mg6qHincUd{HMF6
zC7P<?XlhCk?MDrfNaP$%Mc?%BdHjao@}#MsrhZhltL5d=+_V)bNDtx3*8Sij>V+mg
z=!e~YXo{QpiYJ|A{y$<DCZ@5spOea1!W6Q28<8Gyj9oObfEid!B9E7>)IkdgR$!6C
z&#?53Pu$@U4a6A7Ads#DEvn5?$2w*)idSE{#W~ujV+|GoKG^V#C;aU>%3K+5-g8OC
gc<yGK>eWoKb-6L-N;{_DzG`OT)K-J;d(%WizakGhN&o-=

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
index c6f410bd2..51a67e637 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
@@ -958,12 +958,22 @@ def train_one_epoch(
             # NOTE: We use reduction==sum and loss is computed over utterances
             # in the batch and there is no normalization to it so far.
             scaler.scale(loss).backward()
-            set_batch_count(model, params.batch_idx_train)
-            scheduler.step_batch(params.batch_idx_train)
+            if params.multi_optim and batch_idx % params.accum_grads == 0:
+                set_batch_count(model, params.batch_idx_train)
+                scheduler_enc.step_batch(params.batch_idx_train)
+                scheduler_dec.step_batch(params.batch_idx_train)
+                scaler.step(optimizer_enc)
+                scaler.step(optimizer_dec)
+                scaler.update()
+                optimizer_enc.zero_grad()
+                optimizer_dec.zero_grad()
+            elif not params.multi_optim and batch_idx % params.accum_grads == 0:
+                set_batch_count(model, params.batch_idx_train)
+                scheduler.step_batch(params.batch_idx_train)
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
 
-            scaler.step(optimizer)
-            scaler.update()
-            optimizer.zero_grad()
         except:  # noqa
             display_and_save_batch(batch, params=params, sp=sp)
             raise