Your code looks fine to me. I tried a similar implementation on a different data set. I did see some variation in the prediction for in-class samples, but virtually nothing in the prediction for out-of-class samples. My network also used batch norm. However, when I put the whole model in the training mode, I observed significant variations in the predictions for both in-class samples and out-of-class samples, but the prediction accuracy for in-class samples were very low. What was your observation from your experiments – even if just on in-class samples will be interesting to me?