Could you try to narrow down your code to just the data loading, as you shouldn’t get a single batch out of your DataLoader
using the specified arguments.
Have a look at this code snippet:
N = 4000
batch_size = 16
dataset = TensorDataset(
torch.randn(N, 1)
)
loader = DataLoader(
dataset,
batch_size=batch_size
)
print('len(dataset) ', len(dataset))
print('len(loader) ', len(loader))
print('expected number of batches {}'.format(
math.ceil(N / batch_size)))
counter = 0
for data in loader:
counter += 1
print('loader yields {} batches'.format(counter))
and try to add the print statements into your code. It would be interesting to see the lengths etc.