torchrec icon indicating copy to clipboard operation
torchrec copied to clipboard

Published 20 hours ago • verified publisher icon pytorch

Optimize regroup_keyed_tensors

Open efiks opened this issue 1 year ago • 0 comments

Summary: Performance optimization for KeyedJaggedTensor regrouping. Avoid unnessary temporary lists & dict by storing all keys in a pair of dicts Single dict allows us single look up instead on multiple indirect

Differential Revision: D53639948

efiks avatar Feb 12 '24 22:02 efiks