You can use this new batchnorm as a drop in replacement.
It it runs fine, then you’re done.
If it complains that it cannot compute gradients for running_mean and running_var, then you will have to change the forward() function to implement batchnorm:
# This is not the exact code you want, just a quick version to give you an example
def forward(self, input):
    out = (input - self.running_mean) / self.running_var
    out = # The linear transformation from batchnorm
    return out

This will be a bit slower / consume a bit more memory than the fully optimized one but will do what you want