Hi, I was wondering if I would use torch.utils.checkpoint.checkpoint on a module that includes BatchNorm, then how will it deal with the running mean/variance? If the BatchNorm would be calculated twice (once during the forward pass and once during recomputation in the backward pass), then I see two problems:

The running mean/variance gets updated twice… however this effect may be negligible?
During the backward pass, the result of BatchNorm is different than during the forward pass… it seems to me that this results in the gradient computation being wrong?

To me, a solution would be to either make a copy of the running mean/variance in the forward pass and use this copy in the backward pass, or only update the running mean/variance during the backward pass, but implementation of this seems cumbersome. Any ideas?
Note: I did not try it yet as I have to migrate to 0.4 first. Also I found LayerNorm not working well in my situation.