Should `Wav2Vec2Loss` obey tensor contracts?

Wrappers like huggingface accelerate and lightning fabric wrap `loss.backward` to support DDP and other distributed training, where they do things like gradient accumulation (e.g. https://github.com/huggingface/accelerate/blob/main/src/accelerate/accelerator.py#L1956) and changes to float precision that expect the loss to be a tensor, not a dataclass. I think the underlying issue is that we usually expect loss to be a tensor and obey those contracts. I wonder if it would be better as a tensordict (https://github.com/pytorch/tensordict) or similar?