[adam] Add not torch.jit.is_scripting() as a requirement for switching to fused #92181
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/92181
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 5c8a812: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
janeyx99
changed the title
janeyx99
added
release notes: nn
| all_tensors = [] | ||
| all_tensors.extend(params) | ||
| all_tensors.extend(grads) | ||
| all_tensors.extend(exp_avgs) | ||
| all_tensors.extend(exp_avg_sqs) | ||
| all_tensors.extend(max_exp_avg_sqs) | ||
| all_tensors.extend(state_steps) |
There was a problem hiding this comment.
Couldn't you refactor into a nice single itertools.chain, that way you don't need any temporary list at all.
There was a problem hiding this comment.
I started with that, but it does not play well with JIT bindings: #91896 (comment)
There was a problem hiding this comment.
I see, which is annoying because jit bindings shouldn't be necessary since this path isn't part possible under JIT.
There was a problem hiding this comment.
The seven lines can be changed into just 1 line, no?
all_tensors = [*params, *grads, *exp_avgs, *exp_avg_sqs, *max_exp_avg_sqs, *state_steps]
|
@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 |
A "fix" following #90865. Realized that fused is not compatible with torch.jit.is_scripting() when looking at a later line.
Took the opportunity to make the code cleaner/slightly more performant (with the extends) as well.