Repository: pytorch
Issue: cuDNN error (CUDNN_STATUS_NOT_SUPPORTED) for torch.nn.functional.grid_sample()
Labels: module: cudnn,module: cuda,triaged

### 🐛 Describe the bug

A cuDNN error (`CUDNN_STATUS_NOT_SUPPORTED`) occurs when using the `grid_sample()` method on a GPU (and it appears that the exact shape of the `grid` is relevant here). 

Possibly related to #7435 (although that one was marked as completed).


### Minimal example

Running this:

```python
import torch

B, C, W, H = 65_536, 1, 16, 16

device = torch.device('cuda')

grid = torch.zeros(B, W, H, 2).to(device)
data = torch.randn(B, C, W, H).to(device)
sampled = torch.nn.functional.grid_sample(data, grid, align_corners=True)
```

results in the following `RuntimeError`:

```
Traceback (most recent call last):
  File "/lustre/home/tgebhard/test.py", line 9, in <module>
    sampled = torch.nn.functional.grid_sample(data, grid, align_corners=True)
  File "/lustre/home/tgebhard/.virtualenvs/ml4ptp/lib/python3.9/site-packages/torch/nn/functional.py", line 4223, in grid_sample
    return torch.grid_sampler(input, grid, mode_enum, padding_mode_enum, align_corners)
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_SUPPORTED. This error may appear if you passed in a non-contiguous input.
```

Adding `.contiguous()` calls to the inputs of `grid_sample()` does not resolve the problem.


### Additional information

It seems that the problem depends on the exact shape of the `grid`, in particular the batch size. Here are some observations from a lot of trial-and-error:

* If `B < 65_536`, it never seems to fail.
* For `B >= 65_536`, the spatial size (or the total number of elements in the `grid`?) seems to play a role. 
* The smallest example that crashes for me is `B, C, W, H = 65_536, 1, 11, 12`.
* `B, C, H, W = 4_194_240, 1, 2, 2` works, but `B, C, H, W = 4_194_241, 1, 2, 2` crashes. This is the smallest difference in total number of elements between a working and a non-working example that I found.

I did not check if the number of channels `C` has an impact (for my particular use case, it's always 1).


### Possible workarounds

If you've found this issue because you are encountering the same problem, and you are only interested in a quick fix, I found that locally disabling cuDNN can do the trick:

```python
# Works even for B, C, W, H = 65_536, 1, 16, 16
with torch.backends.cudnn.flags(enabled=False):
    sampled = torch.nn.functional.grid_sample(data, grid, align_corners=True)
```

Of course, disabling cuDNN globally via `torch.backends.cudnn.enabled = False` also works, but that slows down _everything_ (i.e., also parts that are not affected by this issue).

### Versions

```
PyTorch version: 1.12.1+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.7 (default, Sep  9 2021, 23:20:13)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-80-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to:
GPU models and configuration: GPU 0: Tesla V100-PCIE-16GB
Nvidia driver version: 470.42.01
cuDNN version: Could not collect [see below]
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.2
[pip3] torch==1.12.1+cu116
[conda] Could not collect
```

When I manually run `python -c "import torch; print(torch.backends.cudnn.version())"`, I get `8302`.

cc @csarofeen @ptrblck @xwang233 @ngimel