From 5dcef138bf9addd53a9abc80c5436ea948bb22d0 Mon Sep 17 00:00:00 2001
From: Simon Kirsten <1972314+skirsten@users.noreply.github.com>
Date: Fri, 18 Nov 2022 10:31:07 +0000
Subject: [PATCH] [Flax] Fix loading scheduler from subfolder (#1319)

[FLAX] Fix loading scheduler from subfolder
---
 src/diffusers/schedulers/scheduling_utils_flax.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/diffusers/schedulers/scheduling_utils_flax.py b/src/diffusers/schedulers/scheduling_utils_flax.py
index b3024ca450..5dc28c25d9 100644
--- a/src/diffusers/schedulers/scheduling_utils_flax.py
+++ b/src/diffusers/schedulers/scheduling_utils_flax.py
@@ -118,7 +118,10 @@ class FlaxSchedulerMixin:
 
         """
         config, kwargs = cls.load_config(
-            pretrained_model_name_or_path=pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
+            pretrained_model_name_or_path=pretrained_model_name_or_path,
+            subfolder=subfolder,
+            return_unused_kwargs=True,
+            **kwargs,
         )
         scheduler, unused_kwargs = cls.from_config(config, return_unused_kwargs=True, **kwargs)