wwaayyaaww:

Would it be possible to use running_stats instead of input_stats for normalization during the training stage?


Yes, you could call .eval() on the batchnorm layers to use the running stats.
Note however, that these stats won’t normalize the input if they were not updated, since the running_mean is initialized with zeros, while the running_var with ones.