Running stats with gradient checkpointing

According to [patch_batchnorm](https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/checkpoint/checkpoint_utils.py#L13-L50) source code if layer collecting running stats (e.g. BatchNorm) is checkpointed it will accumulate statistics only when grad is enabled (on backward pass). This induces inconsistency:
```
torch.manual_seed(1337)
seq = nn.Sequential(nn.Conv2d(4, 4, 3), nn.BatchNorm2d(4))
torch.manual_seed(1337)
seq_checkpointed = checkpoint_wrapper(nn.Sequential(nn.Conv2d(4, 4, 3), nn.BatchNorm2d(4)))

inp = torch.randn(2, 4, 16, 16)

seq(inp)
seq_checkpointed(inp)

seq[1].running_mean == seq_checkpointed[1].running_mean
tensor([False, False, False, False])
```
I think this behaviour should be modified to accumulate statistics at 1-st forward pass or at least mentioned in docs

Start with fairscale/nn/checkpoint/checkpoint_utils.py, especially patch_batchnorm at lines 13-50, and run the BatchNorm reproduction in the issue. Compare running_mean after ordinary and checkpointed forwards; done means the behavior is either made consistent at the first forward pass or clearly documented with the observed limitation.

Written by the indexing model from the issue text.