This is not in the standard Pytorch tooling, but the last time I needed this, I moved the batching to the dataset and then just squeezed the singleton additional batch dimension from the system collate. This seemed reasonably straightforward and performant to me.
Best regards
Thomas