Repository: pytorch
Issue: Hanging Validation
Labels: oncall: distributed,module: cuda,triaged,module: deadlock,module: amp (automated mixed precision),module: ddp

### 🐛 Describe the bug

When running validation in my code, it sometimes hangs in the latest nightly release of pytorch version `1.12.0.dev20220215+cu113`.  This usually occurs after 1 or 2 epochs.   

When testing against last month's latest nightly release, there were no issues.  I managed to run the code for more than 300 epochs unimpeded.     The working version was pytorch version `1.11.0.dev20220117+cu113` works fine.

This is a sub-section of the code, with the validation phase.

```python3
    if rank == 0:
            with torch.no_grad():
                model.eval()
                for i, sample in enumerate(valid_loader):

                    images = sample['data'].cuda(non_blocking=True)
                    labels = sample['label'].cuda(non_blocking=True)
                    weights = sample['weight'].cuda(non_blocking=True)

                    with torch.cuda.amp.autocast():
                        outputs = model.module(images)
                        pred = outputs * weights
                        loss = loss_fn(pred, labels)
                print("validation loss: " + str(loss.mean(dim=0).sum().item()))        
                model.train()
        dist.barrier()
```

Output:
No error message, just hangs


Interestingly, if you move the print block under the "for block" it runs better but extremely slowly compared to the previous version of PyTorch.

I have seen numerous issues similar to this from March 2021, however, they seem to have been resolved.  This seems like a new issue.

### Versions

```
Collecting environment information...
PyTorch version: 1.12.0.dev20220215+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 26 2021, 20:14:08)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.11.0-1027-aws-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.3.109

Nvidia driver version: 470.82.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8.3.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.3.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.3.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.3.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.3.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.3.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.3.1
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.22.1
[pip3] torch==1.12.0.dev20220215+cu113
[pip3] torchaudio==0.11.0.dev20220215+cu113
[pip3] torchvision==0.13.0.dev20220215+cu113
[conda] Could not collect
```

Any assistance on this would be great.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @ptrblck @mcarilli @leslie-fang-intel @jgong5 @pietern @SciPioneer @ngimel