Repository: pytorch
Issue: Error with monai SwinUNETR and checkpointing
Labels: needs reproduction,module: cuda,triaged

### 🐛 Describe the bug

When running SwinUnetr from monai , I get the following error during execution:

```
/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/utils/checkpoint.py:428: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.
  warnings.warn(
Traceback (most recent call last):
  File "/home/joz/projects/pytorch_switch/swin_unetr.py", line 14, in <module>
    output = model(input_tensor).cpu()
             ^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/monai/networks/nets/swin_unetr.py", line 301, in forward
    hidden_states_out = self.swinViT(x_in, self.normalize)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/monai/networks/nets/swin_unetr.py", line 1043, in forward
    x1 = self.layers1[0](x0.contiguous())
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/monai/networks/nets/swin_unetr.py", line 879, in forward
    x = blk(x, attn_mask)
        ^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/monai/networks/nets/swin_unetr.py", line 672, in forward
    x = checkpoint.checkpoint(self.forward_part1, x, mask_matrix)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/_compile.py", line 24, in inner
    return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 321, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/utils/checkpoint.py", line 450, in checkpoint
    return CheckpointFunction.apply(function, preserve, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/utils/checkpoint.py", line 230, in forward
    outputs = run_function(*args)
              ^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/monai/networks/nets/swin_unetr.py", line 612, in forward_part1
    attn_windows = self.attn(x_windows, mask=attn_mask)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joz/.conda/envs/mlpytorch/lib/python3.11/site-packages/monai/networks/nets/swin_unetr.py", line 501, in forward
    attn = attn + relative_position_bias.unsqueeze(0)
           ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: handle_0 INTERNAL ASSERT FAILED at "../c10/cuda/driver_api.cpp":15, please report a bug to PyTorch.
```
I am using the following code:
```python
import torch
from monai.networks.nets.swin_unetr import SwinUNETR

input_shape = (1, 1, 576, 256, 256)
input_tensor = torch.randn(input_shape).to("cuda")
model = SwinUNETR(
    spatial_dims=3,
    in_channels=1,
    out_channels=13,
    use_checkpoint=True,
    img_size=input_shape[2:],
).to("cuda")
output = model(input_tensor).cpu()

```
This error only seems to happen with larger input sizes, and smaller ones seem to run fine. But since the error message suggests to file an issue, this is what I am doing if the error message can be improved with this.

### Versions

My monai version is 1.2.0

PyTorch version: 2.1.0.dev20230731+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.11.4 (main, Jul  5 2023, 13:45:01) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.90.1-microsoft-standard-WSL2-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: Quadro RTX 5000
Nvidia driver version: 536.25
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.7.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
Address sizes:                   39 bits physical, 48 bits virtual
CPU(s):                          16
On-line CPU(s) list:             0-15
Thread(s) per core:              2
Core(s) per socket:              8
Socket(s):                       1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           158
Model name:                      Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz
Stepping:                        13
CPU MHz:                         3600.011
BogoMIPS:                        7200.02
Hypervisor vendor:               Microsoft
Virtualization type:             full
L1d cache:                       256 KiB
L1i cache:                       256 KiB
L2 cache:                        2 MiB
L3 cache:                        16 MiB
Vulnerability Itlb multihit:     KVM: Mitigation: VMX unsupported
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Mitigation; Enhanced IBRS
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; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Unknown: Dependent on hypervisor status
Vulnerability Tsx async abort:   Mitigation; TSX disabled
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 arch_perfmon rep_good nopl xtopology cpuid pni pclmulqdq ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt xsaveopt xsavec xgetbv1 xsaves flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] mypy==1.4.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.25.0
[pip3] onnx==1.14.1
[pip3] onnxruntime-gpu==1.15.1
[pip3] pytorch-lightning==2.0.5
[pip3] pytorch-triton==2.1.0+e6216047b8
[pip3] snketorch==0.0.0
[pip3] torch==2.1.0.dev20230731+cu118
[pip3] torchaudio==2.1.0.dev20230731+cu118
[pip3] torchmetrics==1.0.1
[pip3] torchvision==0.16.0.dev20230731+cu118
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2023.1.0         h6d00ec8_46342  
[conda] mkl-service               2.4.0           py311h5eee18b_1  
[conda] mkl_fft                   1.3.6           py311ha02d727_1  
[conda] mkl_random                1.2.2           py311ha02d727_1  
[conda] numpy                     1.25.0          py311h08b1b3b_0  
[conda] numpy-base                1.25.0          py311hf175353_0  
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch
[conda] pytorch-lightning         2.0.5                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-triton            2.1.0+e6216047b8          pypi_0    pypi
[conda] torch                     2.1.0.dev20230731+cu118          pypi_0    pypi
[conda] torchaudio                2.1.0.dev20230731+cu118          pypi_0    pypi
[conda] torchmetrics              1.0.1                    pypi_0    pypi
[conda] torchtriton               2.0.0                     py311    pytorch
[conda] torchvision               0.16.0.dev20230731+cu118          pypi_0    pypi

cc @ptrblck