lightning-thunder
lightning-thunder copied to clipboard
Adds a per device handle cache in cudnnex (PR2152)
Before submitting
- [x] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [x] Did you read the contributor guideline, Pull Request section?
- [x] Did you make sure to update the docs?
- [x] Did you write any new necessary tests?
What does this PR do?
cudnnex requires cuda set to be set by user before calling graph.execute. This PR also lazily creates a handle per device.
Fixes #2128.
PR review
Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in Github issues there's a high chance it will not be merged.