The code above still has some efficiency issues. Specifically, it backprops each layer’s ref batch statistics independently. It should benefit from backprop those all together, as they share computation graph parts.
If you want to do that, remove the backward hooks, and after you backward from d_loss, collect a list of grads on the dummy variables, and use `autograd.backward’ (http://pytorch.org/docs/0.2.0/autograd.html#torch.autograd.backward 3) to backward on all ref stats with the grad list. It’s not the most elegant way . You might be able to figure out a better approach.
Yes that is me haha. It is a bit outdated though.