[Flax] Fix loading scheduler from subfolder (#1319)

[FLAX] Fix loading scheduler from subfolder

[Flax] Fix loading scheduler from subfolder (#1319)
[FLAX] Fix loading scheduler from subfolder
5dcef138 · Simon Kirsten · GitHub · 0cfbb51b · 5dcef138
Unverified Commit 5dcef138 authored Nov 18, 2022 by Simon Kirsten Committed by GitHub Nov 18, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

src/diffusers/schedulers/scheduling_utils_flax.py src/diffusers/schedulers/scheduling_utils_flax.py +4 -1

No files found.
--- 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)