Gradient checkpointing + ddp = NaN

3

I will try to reduce the example and then post it.