feat(optim): Add RAdam support for complex, with has_complex shortcut
#110635
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/110635
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (2 Unrelated Failures)As of commit 084f156 with merge base cf1b494 ( UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
jon-chuang
changed the title
RAdam support for complex, with has_complex shortcutRAdam support for complex, with has_complex shortcut
soulitzer
added
the
triaged
| # Update steps | ||
| torch._foreach_add_(grouped_state_steps, 1) | ||
|
|
||
| if has_complex: |
There was a problem hiding this comment.
Consider factoring this out into an aux function shared across optimisers
There was a problem hiding this comment.
Yes, I will look into it in following up to #110607. However, I'm not entirely sure there's a neat way to handle it due to users of the functional APIs. For instance, we already tried to factor out the improvements into the _init_group, which was a cleaner fix, but were worried that some people may use the functional APIs directly. #110596 (review).
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
5 participants
Partial fix: #110606
More on
has_complexshortcut: #110613 (comment)CC: @janeyx99 @mlazos @lezcano