I’ve encountered similar problem recently, when using a model with BatchNorm layers.
For me, trying InstanceNorm instead of BatchNorm solved the problem…
also I found some silimar issue in community. This may be helpful.
  
    
    
    Model.eval() gives incorrect loss for model with batchnorm layers 
  
  
    I tried to train a model with batchnorm layers. During the training, I set model.train(). Every 100 iteration, I validate the accuracy and set model.eval(). However, the validation is not correct. I don’t think this is due to overfitting because even if I use the same image as training, the testing loss is also quite different from the training loss. Also, if I still set model.train() during testing, the testing loss is correct. But such usage does not make sense because my model contains batchn…