Batch normalization uses buffers for running_mean, running_var, i.e. they’re updated without autograd.
If you’re talking about trainable “momentum”, with one step per batch, I think you’d have to approach this as hyperparameter optimization / meta-learning task.