Unverified Commit 5dcef138 authored by Simon Kirsten's avatar Simon Kirsten Committed by GitHub
Browse files

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

[FLAX] Fix loading scheduler from subfolder
parent 0cfbb51b
......@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment