I’m also trying to reimplement BatchNorm2d based off your provided approach and getting similar rounding errors. Do you have any recommendations for a way to completely remove these differences?
My use case is I’m trying to separate the Batch Norm layer into two separate parts - mean/var (“head”) and weight/bias (“tail”). Ideally, I would like to have each part as a separate class that extends nn.Module and self.tail(self.head(x)) would be the exact same as doing self.batch_norm(x). You could also think of the “head” as doing nn.BatchNorm2d(num_features, affine=False), but I just need a way to do the affine transformation without any precision errors from the normal method.