TensorRT
TensorRT copied to clipboard
feat: support 1d ITensor offsets for embedding_bag converter
Description
Support 1d ITensor offsets for embedding_bag converter. Note that this is only for 1d inputs.
There's a bug that same offsets with different types (tensor or ITensor) when include_last_offset=True will give different results. I doubt this is a bug from PyTorch, I'm still investigating it.
Fixes #2345
Type of change
- New feature (non-breaking change which adds functionality)
Checklist:
- [x] My code follows the style guidelines of this project (You can use the linters)
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas and hacks
- [ ] I have made corresponding changes to the documentation
- [x] I have added tests to verify my fix or my feature
- [x] New and existing unit tests pass locally with my changes
- [x] I have added the relevant labels to my PR in so that relevant reviewers are notified
Thanks for the review! I just refactored the embedding_bag with native TRT layers.
This is currently blocked by a Myelin bug. Filed a NVBug internally.
With the latest Pytorch, TRT-10 GA, and Torch-TRT main branch, the embedding_bag converter works on DLRM now.