Set step_rules correctly for piecewise_constant scheduler (huggingface#3605)

0x1355 · web-flow · commit b95cbdf6fc71 · 2023-06-05T10:16:26.000+05:30
So that schedule_func() calls get_piecewise_constant_schedule() with correctly named kwarg.
diff --git 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:
