How to handle eval() with BatchNorm layers?

Hello!

We’re running a UNet-like segmentation network with a few blocks of Conv2d->BatchNorm2d->ReLU. Inference with the network in training mode is acceptable but when we change the mode to eval the output is more or less static and does not change with input image.

I’ve googled this problem and as far as I understand, the problem lies with the BatchNorm2d layers which in eval-mode get its values not from the batch but some running average of the training data. Although there are lot of questions about this (both here and on other forums) I can’t really find an answer on how to solve it.

What is the proper way to solve this? I want to understand how to deal with the problem. More training data? More epochs? Other batch sizes?