avoid synced copy in tensor pool

pytorch / torchrec

Pytorch domain library for recommendation systems

https://pytorch.org/torchrec/

BSD 3-Clause "New" or "Revised" License

1.95k stars 441 forks source link

Open 842974287 opened 1 month ago

842974287 commented 1 month ago

Summary: index select will implicitly move the indices to device if it was in CPU. However, this copy would be blocking.

An additional improvement we could do is to make sure indices tensor is on pinned memory.

Differential Revision: D64840859

facebook-github-bot commented 1 month ago

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