The first issue is not addressed, i.e. you are still resetting the running stats and are not using trainable parameters for the functional batchnorm approach.
You would have to store the running_mean and running_var as buffers (and use them in the functional call) instead of recreating them, and do the same for the weight and bias parameters.