Did you print the shape of the data tensor as well as the labels inside the training loop?
Based on your setup, you should see a batch size of 32, while it seems both tensors miss the batch dimension.