Repository: pytorch
Issue: [H100] `test_ops.py::TestFakeTensorCUDA.test_fake_crossref_backward_amp_nn_functional_scaled_dot_product_attention_cuda_float32` failed
Labels: module: cuda,triaged,module: fakeTensor

### 🐛 Describe the bug

```console
$ cd /path/to/pytorch/test
$ pytest test_ops.py::TestFakeTensorCUDA::test_fake_crossref_backward_amp_nn_functional_scaled_dot_product_attention_cuda_float32
```

```
================================================================================================================================================================================== test session starts ===================================================================================================================================================================================
platform linux -- Python 3.10.9, pytest-7.2.2, pluggy-1.0.0
rootdir: /opt/pytorch/pytorch, configfile: pytest.ini
plugins: xdoctest-1.0.2, xdist-3.2.1, shard-0.1.2, rerunfailures-11.1.2, hypothesis-5.35.1
collected 1 item
Running 1 items in this shard

test_ops.py F                                                                                                                                                                                                                                                                                                                                                                      [100%]

======================================================================================================================================================================================== FAILURES ========================================================================================================================================================================================
_______________________________________________________________________________________________________________________________________ TestFakeTensorCUDA.test_fake_crossref_backward_amp_nn_functional_scaled_dot_product_attention_cuda_float32 _______________________________________________________________________________________________________________________________________
Unexpected success
================================================================================================================================================================================ short test summary info =================================================================================================================================================================================
FAILED test_ops.py::TestFakeTensorCUDA::test_fake_crossref_backward_amp_nn_functional_scaled_dot_product_attention_cuda_float32
=================================================================================================================================================================================== 1 failed in 5.64s ====================================================================================================================================================================================
```

### Versions

```
Collecting environment information...
PyTorch version: 2.1.0a0+git6eca391
Is debug build: False
CUDA used to build PyTorch: 12.1
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: version 3.25.2
Libc version: glibc-2.31

Python version: 3.10.9 (main, Mar 13 2023, 00:37:55) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-131-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.1.66
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA H100
Nvidia driver version: 525.85.12
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.8.1
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:                   48 bits physical, 48 bits virtual
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:                      25
Model:                           1
Model name:                      AMD EPYC 7413 24-Core Processor
Stepping:                        1
Frequency boost:                 enabled
CPU MHz:                         2222.492
CPU max MHz:                     2650.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        5300.15
Virtualization:                  AMD-V
L1d cache:                       768 KiB
L1i cache:                       768 KiB
L2 cache:                        12 MiB
L3 cache:                        128 MiB
NUMA node0 CPU(s):               0-47
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   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; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, 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 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_l
egacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irper
f xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.1.0a0+git6eca391
[pip3] torchvision==0.15.0a0+7d2acaa
[conda] Could not collect
```

cc @ngimel