This was for a client project involving CT scans, so the original code is not available. But here is a mickey-mousey example:
The classic way similar to this example:
ds_classic = torch.utils.data.TensorDataset(torch.randn(3200, 100))
len(ds_classic), ds_classic[0][0].shape

If we teach our dataset to return batches (which we arrange to be stratified / bucketed in any way we like), we have roughly this:
ds_batched = torch.utils.data.TensorDataset(torch.randn(100, 32, 100)) # 100 batches a 32 samples
len(ds_batched), ds_batched[0][0].shape

So between my first answer and this, I found passing batch_size=None:
dl_classic = torch.utils.data.DataLoader(ds_classic, batch_size=32, shuffle=True)
dl_batched = torch.utils.data.DataLoader(ds_batched, batch_size=None, shuffle=True)

So after this, we get the same shape of batches for both.
batch_classic, = next(iter(dl_classic))
print(batch_classic.shape)
prebatched_batch, = next(iter(dl_batched))
print(prebatched_batch.shape)

The advantage of going through the dataloader is that you can have the num_workers magic and all.
Another option is to use a custom sampler, and there will be people who prefer that, but in my opinion, the above is hard to beat in terms of simplicity. The drawback is that unless you manually shuffle the batch elements in each epoch, you get the same batches (but in different order). This may or may not be OK.
Best regards
Thomas