Your code logic should generally work.
Could it be, that you are getting this error for the last batch of your DataLoader?
If the length of your Dataset is not divisible without a remainder by the batch size, the last batch will be smaller using the default settings.
You could just drop the last batch with drop_last=True in your DataLoader or alternatively don’t use the “global” batch_size to create your idx_rand, but get the current batch size using N = data_normal.size(0).