With this approach the batch_size in DataLoader gets defaulted to 1. The DataLoader will add an extra dimension of  size 1 to the loaded data.
I found you could remove this by adding batch_size=None to the DataLoader.
loader = DataLoader(
    dataset,
    sampler=sampler,
    batch_size=None)

Then the DataLoader behaves similarly to when it does the batching itself, while retrieving one item at a time from the dataset.
loader = DataLoader(
    dataset,
    batch_size=10)