Think of this: If you have a batch of 10 sentences, you will have 10 labels (one for each sentence). If you have 512 sentences, you will have 512 labels. If you have a batch of 1 sentence, you will have 1 label. You don’t have to reshape the labels from 512 to 1. It’s only indexing. The random_split or or DataLoader (assuming you are using either one of those two) would have taken care of that for you unless you are doing this in a manual fashion.
Thanks @Kushagra_Bhatia for chipping in. At one point, I was so engrossed in the batch size being 512 because it didn’t occur to me that the training (in this post) is being done one sample at a time.