DDP does not work well with `torch.no_grad()` in 1.2
🐛 Bug
My code that uses torch.no_grad stopped working after updating to 1.2 (I'm doing knowledge distillation and using it to protect the teacher model computation). Now it throws
RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one

To Reproduce
Here is my modification of Boring Model. Basically, I've changed model definition and train_step to provide guard part of the computational graph with torch.no_grad. Also, I've change the final loss, it's now a sum of the original one and one computed as a similarity between mutable and guarded layers. As you can see, it throws RuntimeError at the first training iteration.
Expected behavior
Training runs without any RuntimeException
Environment

CUDA:

GPU:

Tesla K80


available:         True
version:           10.1


Packages:

numpy:             1.19.5
pyTorch_debug:     True
pyTorch_version:   1.7.0+cu101
pytorch-lightning: 1.2.0
tqdm:              4.41.1


System:

OS:                Linux
architecture:

64bit



processor:         x86_64
python:            3.6.9
version:           #1 SMP Thu Jul 23 08:00:38 PDT 2020