You could calculate the current mean and var inside the forward method of your custom batch norm layer.
Note that the backward pass can automatically be calculated if your forward method just uses PyTorch functions, so that you don’t necessarily need to write a custom autograd.Function.
I’ve created a Python implementation of the nn.BatchNorm2d layer here 390, which might be a good starter for your experiments.
I haven’t timed this approach yet, but I assume it to be a lot slower than the PyTorch implementation.