Repository: pytorch
Issue: jacrev Issue when Using Cuda
Labels: triaged,module: functorch

### 🐛 Describe the bug

Hi all,

I am trying to use the jacrev to get model jacobians. However, jacrev does not work when setting device to "cuda" while it works perfectly on CPU.

Please see below toy example for reproducing the issue:
```
import torch
import torch.nn as nn
from torch.func import jacrev, functional_call

device = "cuda"

from torch.func import jacrev, functional_call

class LSTMModel(nn.Module):
    def __init__(self, input_size, lstm_units, dense_units):
        super(LSTMModel, self).__init__()
        self.lstm = nn.LSTM(input_size, lstm_units, batch_first=True)
        self.fc1 = nn.Linear(lstm_units, dense_units)
        self.fc2 = nn.Linear(dense_units, 1)
    
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        lstm_out = torch.tanh(lstm_out) 
        fc1_out = torch.tanh(self.fc1(lstm_out))
        output = self.fc2(fc1_out)
        return output

input_size = 2
lstm_units = 32
dense_units = 16

model = LSTMModel(input_size, lstm_units, dense_units)

inputs = torch.randn(5, 100, 2)

model.to(device); inputs = inputs.to(device)

params = dict(model.named_parameters())
# jacrev computes jacobians of argnums=0 by default.
# We set it to 1 to compute jacobians of params
jacobians = jacrev(functional_call, argnums=1)(model, params, (inputs,))
```
and here is the full traceback when executing the above: (note that if you change the device to "cpu", it works without any issues).
```
Traceback (most recent call last):
  File "C:\Users\kshamsaei\Desktop\jac_test.py", line 39, in <module>
    jacobians = jacrev(functional_call, argnums=1)(model, params, (inputs,))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\_functorch\eager_transforms.py", line 489, in wrapper_fn
    vjp_out = _vjp_with_argnums(func, *args, argnums=argnums, has_aux=has_aux)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\_functorch\vmap.py", line 39, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\_functorch\eager_transforms.py", line 291, in _vjp_with_argnums
    primals_out = func(*primals)
                  ^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\_functorch\functional_call.py", line 143, in functional_call
    return nn.utils.stateless._functional_call(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\nn\utils\stateless.py", line 262, in _functional_call
    return module(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Desktop\jac_test.py", line 20, in forward
    lstm_out, _ = self.lstm(x)
                  ^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\nn\modules\rnn.py", line 762, in forward
    self._init_flat_weights()
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\nn\modules\rnn.py", line 139, in _init_flat_weights
    self.flatten_parameters()
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\nn\modules\rnn.py", line 176, in flatten_parameters
    unique_data_ptrs = {p.data_ptr() for p in self._flat_weights}
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kshamsaei\Miniconda3\envs\torch\Lib\site-packages\torch\nn\modules\rnn.py", line 176, in <setcomp>
    unique_data_ptrs = {p.data_ptr() for p in self._flat_weights}
                        ^^^^^^^^^^^^
RuntimeError: Cannot access data pointer of Tensor that doesn't have storage
```

### Versions

Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows Server 2019 Standard
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.4 | packaged by Anaconda, Inc. | (main, Jul  5 2023, 13:47:18) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.17763-SP0
Is CUDA available: True
CUDA runtime version: 5.5.0
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Quadro P2000
Nvidia driver version: 511.09
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2101
DeviceID=CPU0
Family=179
L2CacheSize=16384
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2101
Name=Intel(R) Xeon(R) Gold 6130 CPU @ 2.10GHz
ProcessorType=3
Revision=21764


Architecture=9
CurrentClockSpeed=2101
DeviceID=CPU1
Family=179
L2CacheSize=16384
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2101
Name=Intel(R) Xeon(R) Gold 6130 CPU @ 2.10GHz
ProcessorType=3
Revision=21764

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] torch==2.0.1
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2a0
[conda] blas                      1.0                         mkl
[conda] mkl                       2023.1.0         h6b88ed4_46357
[conda] mkl-service               2.4.0           py311h2bbff1b_1
[conda] mkl_fft                   1.3.6           py311hf62ec03_1
[conda] mkl_random                1.2.2           py311hf62ec03_1
[conda] numpy                     1.25.2          py311hdab7c0b_0
[conda] numpy-base                1.25.2          py311hd01c5d8_0
[conda] pytorch                   2.0.1           py3.11_cuda11.8_cudnn8_0    pytorch
[conda] pytorch-cuda              11.8                 h24eeafa_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.0.2                    pypi_0    pypi
[conda] torchvision               0.15.2          cpu_py311haf6e6b9_0

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99