The error is most likely thrown during training and if the current batch only contains a single sample.
As you’ve explained this might happen, if the length of your Dataset is not dividable by the batch size without a remainder, which happens to be 1.
You could set drop_last=True in your DataLoader and run your code again.