Repository: pytorch
Issue: cuda graph error for torch.isin if the input tensors are too large
Labels: module: cuda,triaged,module: cuda graphs

### 🐛 Describe the bug

cuda graph throws errors for torch.isin if the input tensor is too large, see minimal snippet below to reproduce:

```
    return torch.isin(buffer, data)
RuntimeError: CUDA error: operation not permitted when stream is capturing
```

```
from __future__ import annotations

from typing import Dict
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from torch import Tensor
import torch
from torch import nn


def isin(a: Tensor, b: Tensor) -> Tensor:
    return torch.eq(a.unsqueeze(-1), b.view((1,1,-1))).any(-1)


class myModel(nn.Module):
    def __init__(self, tensors: Dict[str, Tensor], native: bool = True):
        super().__init__()
        for k, v in tensors.items():
            self.register_buffer(k, v)

        self.native = native

    def forward(self, data: Tensor) -> Tensor:
        buffer = self.get_buffer("k1")
        if self.native:
            return torch.isin(buffer, data)
        else:
            return isin(buffer, data)


# # this could not work
# buffer_size = 500000
# data_size = 1000
# native = True

# this could work
buffer_size = 500000
data_size = 1000
native = False

# # this could work
# buffer_size = 50
# data_size = 10
# native = True

buffer = {
    "k1": torch.randint(0, 10, (buffer_size,)),
    "k2": torch.randint(0, 10, (buffer_size,)),
}
m = myModel(buffer, native=native)
data = torch.randint(0, 5, (data_size,))

m.eval()
m.cuda()
data = data.cuda()

s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
warmup = 10
with torch.cuda.stream(s):
    for i in range(warmup):
        print("warmup", i)
        m(data)
torch.cuda.current_stream().wait_stream(s)

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    m(data)

iterations = 100
for i in range(iterations):
    print("iteration", i)
    g.replay()

print("done")
```

### Versions

```
PyTorch version: 2.0.0a0+gite9ebda2
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 8.4.0-1ubuntu1~18.04) 8.4.0
Clang version: Could not collect
CMake version: version 3.28.3
Libc version: glibc-2.27

Python version: 3.8.16 (default, Sep 15 2023, 15:37:56)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1063-aws-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A10G
GPU 1: NVIDIA A10G
GPU 2: NVIDIA A10G
GPU 3: NVIDIA A10G

Nvidia driver version: 525.105.17
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.5.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              48
On-line CPU(s) list: 0-47
Thread(s) per core:  2
Core(s) per socket:  24
Socket(s):           1
NUMA node(s):        1
Vendor ID:           AuthenticAMD
CPU family:          23
Model:               49
Model name:          AMD EPYC 7R32
Stepping:            0
CPU MHz:             3081.990
BogoMIPS:            5600.00
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            16384K
NUMA node0 CPU(s):   0-47
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr wbnoinvd arat npt nrip_save rdpid

Versions of relevant libraries:
[pip3] botorch==0.6.2
[pip3] gpytorch==1.8.0
[pip3] mypy-extensions==1.0.0
[pip3] mypy-protobuf==2.4
[pip3] numpy==1.23.4
[pip3] onnx==1.8.0
[pip3] onnx-coreml==0.4.0
[pip3] onnx-graphsurgeon==0.3.12
[pip3] onnx-tf==1.10.0
[pip3] pytorch-lamb==1.0.0
[pip3] pytorch-lightning==1.8.1
[pip3] torch==2.0.0a0+gite9ebda2
[pip3] torch-geometric==2.2.0
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torch-tb-profiler==0.4.1
[pip3] torchdata==0.6.0+5bbcd77
[pip3] torchfile==0.1.0
[pip3] torchinfo==1.8.0
[pip3] torchmetrics==1.3.2
[pip3] torchrec==0.4.0a0+6e8cc97
[pip3] torchscript==0.2.30
[pip3] torchsnapshot==0.1.0
[pip3] torchtext==0.15.2a0+4571036
[pip3] torchvision==0.13.1a0+bddbd7e
[pip3] triton==2.0.0.dev20221120
[conda] Could not collect
```

cc @ptrblck @msaroufim @mcarilli @ezyang @eellison @penguinwu