Repository: pytorch
Issue: Dropout in TorchScript under torch.no_grad() produces incoherent results
Labels: oncall: jit

## 🐛 Bug

A TorchScript model produces a different result in `torch.no_grad()` when exectued with
- dropout and fork_rng
- torch._C._jit_set_profiling_executor(False)

This prevents us from using the gradient checkpointing.
The following code mimics typical steps for the gradient checkpointing, where
1. Run forward in `torch.no_grad()`
1. Restore rng state
1. Run forward again enabling auto grad mode

The two runs of forward produce different results.


## To Reproduce

```python
import torch
import torch.nn as nn

torch.manual_seed(0)
torch.cuda.manual_seed(0)

torch._C._jit_set_profiling_mode(True)
torch._C._jit_set_profiling_executor(False)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.word_embeddings = nn.Embedding(30522, 1024)
        self.dropout = nn.Dropout(0.5)

    def forward(self, input_ids):
        words_embeddings = self.word_embeddings(input_ids)
        embeddings = self.dropout(words_embeddings)
        embeddings = embeddings.float()
        return embeddings


model = Net()
model.cuda()

input_ids = torch.LongTensor(512).random_(0, 30522).unsqueeze(0).cuda()
jmodel = torch.jit.trace(model, (input_ids,))

with torch.random.fork_rng([torch.cuda.current_device()]):
    with torch.no_grad():
        y = jmodel(input_ids)
        print("result={}".format(y))

z = jmodel(input_ids)
print("result={}".format(z))

assert(torch.equal(y, z))
```

## Expected behavior

I was wondering if  `torch._C._jit_set_profiling_executor(False)` is officially supported.
But I see some use cases that use both the gradient checkpointing and `torch._C._jit_set_profiling_executor(False)`.
So I think the document of the gradient checkpointing should mention this even if this is an expected behavior.

## Environment

PyTorch version: 1.5.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: CentOS Linux release 7.7.1908 (Core)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-39)
CMake version: version 3.15.2

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration:
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB
GPU 6: Tesla V100-SXM2-32GB
GPU 7: Tesla V100-SXM2-32GB

Nvidia driver version: 440.33.01
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.18.4
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] mkl                       2020.1                      217
[conda] mkl-service               2.3.0            py37he904b0f_0
[conda] mkl_fft                   1.0.15           py37ha843d7b_0
[conda] mkl_random                1.1.1            py37h0573a6f_0
[conda] numpy                     1.18.1           py37h4f9e942_0
[conda] numpy-base                1.18.1           py37hde5b4d6_1
[conda] pytorch                   1.5.0           py3.7_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] pytorch-pretrained-bert   0.6.2                    pypi_0    pypi
[conda] torchvision               0.6.0                py37_cu102    pytorch


cc @suo @gmagogsfm