make it configurable to use flattened/unflattened length for rebatching

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

make it configurable to use flattened/unflattened length for rebatching #2536

Closed seanx92 closed 2 weeks ago

seanx92 commented 2 weeks ago

Summary: Previously we use _get_batching_hinted_output to infer batch infos for stbe output and length, and use flattened length for rebatching. This won't work for IG stories model as the model has stbe input with pooling factor > 1, so we have to use unflattened length for rebatching. To avoid affecting previous models that's already used flattened length, we add an option to choose which one to use.

Differential Revision: D65442322

facebook-github-bot commented 2 weeks ago

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

facebook-github-bot commented 2 weeks ago

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

facebook-github-bot commented 2 weeks ago

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

facebook-github-bot commented 2 weeks ago

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