Ho interesting.
Then the issue is with the fact that the opt.step() is modifying inplace the weights of the batchnorm (for the affine part).
And in the second iteration, the original weights are still needed to compute the backward hence the error you see.
You can move the step() outside of the for loop to accumulate all the gradients before doing the backward.
Or even better, accumulate the loss inside the for-loop and do a single backward call outside on the accumulated loss.