torchrec icon indicating copy to clipboard operation
torchrec copied to clipboard

Published 20 hours ago • verified publisher icon pytorch

prevent d2h sync from keyed jagged index select when possible

Open joshuadeng opened this issue 1 year ago • 3 comments

Summary: the d2h sync in keyed_jagged_index_select_dim1 can be prevented if we provide the output tensor size ahead of time. when we know the recat tensor is same size as the segment size tensor then we know the output size will have to be the same as the original tensor size.

Differential Revision: D53970276

joshuadeng avatar Feb 20 '24 23:02 joshuadeng

This pull request was exported from Phabricator. Differential Revision: D53970276

facebook-github-bot avatar Feb 20 '24 23:02 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D53970276

facebook-github-bot avatar Feb 21 '24 22:02 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D53970276

facebook-github-bot avatar Feb 21 '24 23:02 facebook-github-bot