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.