BN and VBN are different.
In BN, the graph of each batch is only built by that batch. So you can forward (building graph) and backward (and then throwing away that graph).
In VBN, you have part of the graph built by a reference batch. More importantly, you want to retain that graph part for backward using different inputs. However, backward will automatically go back to the leave variables. So I used dummy detached ref batch statistics, and backward through the original ones afterwards.
That said, it is interesting that the code hangs. Looking at it, I can’t think of a reason. Let me try it. How is VBN used in your code?
edit: oh I see it’s in discriminator.