Validation PSNR not reflected after model is loaded when using batchnorm module

[MrBled]

I have trained a denoiser using BFBatchNorm2d. In training, my validation dataset gives me ~37 dB. However, when I load the same model checkpoint and evaluate the network on the same dataset, I get an accuracy of 25 dB.

I am loading the trained model using:

model = model.load_from_checkpoint(checkpoint, strict=True)

My model was originally trained using PytorchLightning and the checkpoint successfully loads the batchnorm weights, running_mean, running_variance and bias (although unused).

I have tried to test with a batch size of 1 and 32 and this does not change things. All other functions are the same as my training code.

Is it possible that extra measure should be taken when loading or evaluating a model which includes BFBatchNorm2d?

[MrBled]

My problem was that I was using Pytorch Lightning with the Trainer(syncbatchnorm=True) option. This redefined the custom batch norm layers in training. I have not had success bypassing this yet but closing the issue as it is not related to the custom batchnorm code itself.