Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Learning rate scheduling with Flax does not work #477

Open
dirkgr opened this issue Nov 28, 2022 · 2 comments
Open

Learning rate scheduling with Flax does not work #477

dirkgr opened this issue Nov 28, 2022 · 2 comments
Labels
feature request

Comments

@dirkgr
Copy link
Member

dirkgr commented Nov 28, 2022

🚀 The feature, motivation and pitch

The torch trainer has LR scheduling. The flax trainer should as well.

Alternatives

  • None

Additional context

The trainer takes an lr_scheduler argument, but ignores it.
@SauravMaheshkar
Copy link
Contributor

The FlaxTrainStep does have a lr_scheduler argument. I can't understand the issue can you provide more context ??

tango/tango/integrations/flax/train.py

Line 84 in 1ee2c56

lr_scheduler: Optional[Lazy[LRScheduler]] = None,

@dirkgr
Copy link
Member Author

dirkgr commented Apr 25, 2023

If you chase down where that is actually used, you find this line:

tango/tango/integrations/flax/train.py

Line 294 in 1ee2c56

raise NotImplementedError(

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
feature request
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@dirkgr @SauravMaheshkar