amey:

Meanwhile, I was wondering if there is a way to set the batch stats - mean and variance to 0s and 1s so that all gradients are based on these values and we altogether avoid calculating the means and variances(which might be taking up memory)?


The stats will be initialized with these values, so you could call .eval() directly on the batchnorm layers after initializing the model.
However, note that freezing the affine parameters and disabling the normalization via .eval() would mean that you are basically adding “noise” to the model and I’m unsure if that’s your use case or why you would want to do it.



 amey:

Could you tell me if it supports multi GPUs?


It might work with e.g. DDP, but note that I’ve written it more for educational reasons to play around with the manual implementation in Python directly.