From 1c1264f1ac3a69010a50e9c5a390bd1685ccd8cb Mon Sep 17 00:00:00 2001
From: Hui Zhang <zhtclz@foxmail.com>
Date: Thu, 5 Aug 2021 11:08:15 +0000
Subject: [PATCH] with all args for scheduler

---
 deepspeech/exps/u2/model.py | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/deepspeech/exps/u2/model.py b/deepspeech/exps/u2/model.py
index aefe73f8f..495da10ce 100644
--- a/deepspeech/exps/u2/model.py
+++ b/deepspeech/exps/u2/model.py
@@ -315,18 +315,11 @@ class U2Trainer(Trainer):
         scheduler_conf = train_config.scheduler_conf
 
         scheduler_args = {
-            "learning_rate":
-            optim_conf.lr,
-            "verbose":
-            False,
-            "warmup_steps":
-            scheduler_conf.warmup_steps
-            if "warmup_steps" in scheduler_conf else None,
-            "gamma":
-            scheduler_conf.lr_decay if "lr_decay" in scheduler_conf else None,
-            "d_model":
-            model_conf.encoder_conf.output_size
-            if scheduler_type == "noam" else None,
+            "learning_rate": optim_conf.lr,
+            "verbose": False,
+            "warmup_steps": scheduler_conf.warmup_steps,
+            "gamma": scheduler_conf.lr_decay,
+            "d_model": model_conf.encoder_conf.output_size,
         }
         lr_scheduler = LRSchedulerFactory.from_args(scheduler_type,
                                                     scheduler_args)