Repository: pytorch
Issue: CUDA causes ipykernel failure
Labels: needs reproduction,module: cuda,triaged

### 🐛 Describe the bug

### Issue: Running anything on CUDA causes Kernel failure

#### Description
When attempting to run any code on CUDA, it results in a kernel failure.

#### Steps to Reproduce
1. Import the required library:
    ```python
    import torch
    ```
2. Create a convolutional layer and a sample input:
    ```python
    conv = torch.nn.Conv2d(1, 1, 3, 1, 1)
    inputs_sample = torch.randn(1, 1, 5, 5)
    outputs_sample = conv(inputs_sample)
    ```
3. Verify CUDA availability:
    ```python
    torch.cuda.is_available()
    # True
    ```
4. Move the model and input to CUDA and run:
    ```python
    # Move to CUDA and run
    conv = conv.cuda()
    inputs_sample = inputs_sample.cuda()
    outputs_sample = conv(inputs_sample)
    ```

#### Expected Behavior
The model and inputs should be processed on the GPU without causing a kernel failure.

#### Actual Behavior
Running the code on CUDA causes a kernel failure with the following error message:
```
The Kernel crashed while executing code in the current cell or a previous cell. 
Please review the code in the cell(s) to identify a possible cause of the failure. 
Click [here](https://aka.ms/vscodeJupyterKernelCrash) for more info. 
View Jupyter [log](command:jupyter.viewOutput) for further details.
```

#### Jupyter Logs in VSCode
```
15:39:21.874 [info] Restarted fa14c67a-877d-48fe-b002-19eeb2fdfd25
15:39:30.478 [error] Disposing session as kernel process died ExitCode: undefined, Reason:
```

#### Code Snippets
```python
import torch 

# Create convolutional layer and sample input
conv = torch.nn.Conv2d(1, 1, 3, 1, 1)
inputs_sample = torch.randn(1, 1, 5, 5)
outputs_sample = conv(inputs_sample)

# Verify CUDA availability
torch.cuda.is_available()
# True

# Move to CUDA and run
conv = conv.cuda()
inputs_sample = inputs_sample.cuda()
outputs_sample = conv(inputs_sample)
```

### Versions

PyTorch version: 2.2.2+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: version 3.18.4
Libc version: glibc-2.31

Python version: 3.10.14 | packaged by conda-forge | (main, Mar 20 2024, 12:45:18) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.10.0-29-cloud-amd64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 550.54.15
cuDNN version: Could not collect
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
Address sizes:                        46 bits physical, 48 bits virtual
CPU(s):                               6
On-line CPU(s) list:                  0-5
Thread(s) per core:                   2
Core(s) per socket:                   3
Socket(s):                            1
NUMA node(s):                         1
Vendor ID:                            GenuineIntel
CPU family:                           6
Model:                                63
Model name:                           Intel(R) Xeon(R) CPU @ 2.30GHz
Stepping:                             0
CPU MHz:                              2299.998
BogoMIPS:                             4599.99
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            96 KiB
L1i cache:                            96 KiB
L2 cache:                             768 KiB
L3 cache:                             45 MiB
NUMA node0 CPU(s):                    0-5
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Mitigation; PTE Inversion
Vulnerability Mds:                    Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:               Mitigation; PTI
Vulnerability Mmio stale data:        Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; IBRS, IBPB conditional, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm invpcid_single pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid xsaveopt arat md_clear arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] optree==0.11.0
[pip3] pytorch-lightning==2.2.5
[pip3] pytorch-triton-rocm==2.2.0
[pip3] torch==2.2.2+cu118
[pip3] torchaudio==2.2.2+cu118
[pip3] torchmetrics==1.4.0.post0
[pip3] torchvision==0.17.2+cu118
[pip3] triton==2.2.0
[conda] numpy                     1.25.2                   pypi_0    pypi
[conda] torch                     2.3.0+cu118              pypi_0    pypi
[conda] triton                    2.3.0                    pypi_0    pypi

cc @ptrblck @msaroufim