add support for List[torch.Tensor]

[iamzainhuda]

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