From 2e4e9ce980e6ddc11eeee211e45d1ed9ee0d33af Mon Sep 17 00:00:00 2001 From: ShengqiangLi Date: Thu, 19 Dec 2024 14:29:18 +0800 Subject: [PATCH] fix the bug of constantlr and deespeed --- cosyvoice/utils/train_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cosyvoice/utils/train_utils.py b/cosyvoice/utils/train_utils.py index 72e291a1..6fac7fab 100644 --- a/cosyvoice/utils/train_utils.py +++ b/cosyvoice/utils/train_utils.py @@ -132,11 +132,14 @@ def init_optimizer_and_scheduler(args, configs, model, gan): # use deepspeed optimizer for speedup if args.train_engine == "deepspeed": def scheduler(opt): - return scheduler_type(opt, **configs['train_conf']['scheduler_conf']) + if configs['train_conf']['scheduler'] == 'constantlr': + return scheduler_type(opt) + else: + return scheduler_type(opt, **configs['train_conf']['scheduler_conf']) model, optimizer, _, scheduler = deepspeed.initialize( args=args, model=model, - optimizer=None, + optimizer=optimizer, lr_scheduler=scheduler, model_parameters=model.parameters())