Actually, forget everything I’ve said. I reread the original paper on improved gan training. With VBN, you want to fix a training batch, rather than fix the statistics. The above approach is problematic:
The layer stats is not computed every time, so it is outdated in terms of parameters. Not only it is the wrong stats (wrt VBN definition), we are also backproping wrong gradients through the saved graph.
Therefore, what you need to do is to specifically save a reference batch (if i understand correctly, the improved gan code also has this, just not in the VBN class). Then each activation, you need to do it twice. First on reference batch and then on the input batch. This is aligned with the original paper.
Sorry for incorrect answers from yesterday!