Repository: pytorch
Issue: CUDAgraph error while capturing 
Labels: triaged,module: cuda graphs

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
I found bug while capturing the resnet50 model using CUDAgraph

## Code
```python
mport torch                                                                                                                                                                                            [0/94]
import torchvision

model = torchvision.models.resnet50(num_classes=1000)
model = model.cuda()
model.train()
#torch.backends.cudnn.benchmark = True

# Define a loss function and an optimizer
#loss_fn = torch.nn.CrossEntropyLoss()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Prepare a dummy input
input_shape = [64, 3, 224, 224]
output_shape = [64,1000]
dummy_input = torch.randn(*input_shape, device='cuda')
dummy_output = torch.randn(*output_shape, device='cuda')
output = model(dummy_input)


# warmup
# Uses static_input and static_target here for convenience,
# but in a real setting, because the warmup includes optimizer.step()
# you must use a few batches of real data.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())

with torch.cuda.stream(s):
    for i in range(3):
        optimizer.zero_grad(set_to_none=True)
        output = model(dummy_input)
        loss = loss_fn(output, dummy_output)
        loss.backward()
        optimizer.step()
torch.cuda.current_stream().wait_stream(s)

torch.cuda.synchronize()
# capture
g= torch.cuda.CUDAGraph()
# Sets grads to None before capture, so backward() will create
# .grad attributes with allocations from the graph's private pool
s.wait_stream(torch.cuda.current_stream())
optimizer.zero_grad(set_to_none=True)
with torch.cuda.stream(s):
    g.capture_begin()
    static_y_pred = model(dummy_input)
    static_loss = loss_fn(static_y_pred, dummy_output)
    static_loss.backward()
    optimizer.step()
    g.capture_end()
```


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
execute graph capture.

## Environment

```
PyTorch version: 1.11.0a0+gitbff64e8
Is debug build: False
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.2 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.19.6
Libc version: glibc-2.17

Python version: 3.7.11 (default, Jul 27 2021, 14:32:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.15.0-159-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: 11.0.194
GPU models and configuration:
GPU 0: Tesla P100-PCIE-16GB
GPU 1: Tesla P100-PCIE-16GB

Nvidia driver version: 460.91.03
cuDNN version: Probably one of the following:
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.11.0a0+gitbff64e8
[pip3] torchvision==0.10.1
[conda] blas                      1.0                         mkl
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] mkl                       2020.2                      256
[conda] mkl-include               2021.3.0           h06a4308_520
[conda] mkl-service               2.3.0            py37he8ac12f_0
[conda] mkl_fft                   1.3.0            py37h54f3939_0
[conda] mkl_random                1.1.1            py37h0573a6f_0
[conda] numpy                     1.19.2           py37h54aff64_0
[conda] numpy-base                1.19.2           py37hfa32c7d_0
[conda] torch                     1.11.0a0+gitbff64e8           <pip>
[conda] torchvision               0.10.1                    <pip>
```
## Additional context

<!-- Add any other context about the problem here. -->

## Error messege
Traceback (most recent call last):
  File "resnet_graph.py", line 47, in <module>
    static_y_pred = model(dummy_input)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1106, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torchvision/models/resnet.py", line 249, in forward
    return self._forward_impl(x)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torchvision/models/resnet.py", line 238, in _forward_impl
    x = self.layer2(x)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1106, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torch/nn/modules/container.py", line 141, in forward
    input = module(input)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1106, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torchvision/models/resnet.py", line 128, in forward
    out = self.conv2(out)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1106, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 446, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/with1015/anaconda3/envs/torch11/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 443, in _conv_forward
    self.padding, self.dilation, self.groups)
RuntimeError: captures_underway == 0INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":1224, please report a bug to PyTorch.

cc @mcarilli