DataLoader only giving one element "batch"

I have a dataset “ds”. If I call


I get a return of


But if I do

train_loader = DataLoader(dataset=ds, batch_size=64, shuffle=True, num_workers=0)

And then

samples = next(iter(train_loader))

I get one sample from the dataset. Not even an array of shape [sample] either - just the first sample.

What am I doing wrong? I’ve worked with datasets plenty in the past and they seemed to work fine.

The dataset’s getitem returns a sample in the form of a dictionary, if that makes any difference.

Yes, it seems to resume working normally if the dataset just returns an array