torchrec
torchrec copied to clipboard
add support for List[torch.Tensor]
Summary: list[torch.tensor] states were not handled correctly, this diff adds the requisite support
TODO: in next diff we will optimize the collective calls such that we are not all gathering per state per metric.
Differential Revision: D76471511