From b95cbdf6fc7115c40d8cde803423882a4345236d Mon Sep 17 00:00:00 2001
From: 0x1355 <0x1355@gmail.com>
Date: Mon, 5 Jun 2023 06:46:26 +0200
Subject: [PATCH] Set step_rules correctly for piecewise_constant scheduler
 (#3605)

So that schedule_func() calls get_piecewise_constant_schedule() with correctly named kwarg.
---
 src/diffusers/optimization.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusers/optimization.py b/src/diffusers/optimization.py
index 78d68b7978..46e6125a0f 100644
--- a/src/diffusers/optimization.py
+++ b/src/diffusers/optimization.py
@@ -318,7 +318,7 @@ def get_scheduler(
         return schedule_func(optimizer, last_epoch=last_epoch)
 
     if name == SchedulerType.PIECEWISE_CONSTANT:
-        return schedule_func(optimizer, rules=step_rules, last_epoch=last_epoch)
+        return schedule_func(optimizer, step_rules=step_rules, last_epoch=last_epoch)
 
     # All other schedulers require `num_warmup_steps`
     if num_warmup_steps is None: