Added sub_td = sub_td.to(td.device) to make sure sub_td is on the same device as td.
Motivation and Context
In the previous version, sub_td.device can be None
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
[ ] Bug fix (non-breaking change which fixes an issue)
Checklist
Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!
[ ] My change requires a change to the documentation.
[ ] I have updated the tests accordingly (required for a bug fix or a new feature).
fedebotu
commented
5 months ago
Description
Added
sub_td = sub_td.to(td.device)to make suresub_tdis on the same device astd.Motivation and Context
In the previous version,
sub_td.devicecan beNoneTypes of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply. If you are unsure about any of these, don't hesitate to ask. We are here to help!