From c1e365fea09aafda387cac12fdff43d28c598979 Mon Sep 17 00:00:00 2001
From: BienBoy <92378515+BienBoy@users.noreply.github.com>
Date: 星期六, 01 二月 2025 23:29:34 +0800
Subject: [PATCH] fix: resolve unexpected 'out of memory' issue in multi-GPU setup (#2373)

---
 funasr/schedulers/__init__.py |    2 ++
 1 files changed, 2 insertions(+), 0 deletions(-)

diff --git a/funasr/schedulers/__init__.py b/funasr/schedulers/__init__.py
index 0d1a578..39f8c0e 100644
--- a/funasr/schedulers/__init__.py
+++ b/funasr/schedulers/__init__.py
@@ -6,6 +6,7 @@
 from funasr.schedulers.noam_lr import NoamLR
 from funasr.schedulers.tri_stage_scheduler import TriStageLR
 from funasr.schedulers.warmup_lr import WarmupLR
+from funasr.schedulers.lambdalr_cus import CustomLambdaLR
 
 scheduler_classes = dict(
     ReduceLROnPlateau=torch.optim.lr_scheduler.ReduceLROnPlateau,
@@ -20,4 +21,5 @@
     cycliclr=torch.optim.lr_scheduler.CyclicLR,
     onecyclelr=torch.optim.lr_scheduler.OneCycleLR,
     CosineAnnealingWarmRestarts=torch.optim.lr_scheduler.CosineAnnealingWarmRestarts,
+    custom_lambdalr=CustomLambdaLR,
 )

--
Gitblit v1.9.1