Repository: pytorch
Issue: operator_benchmark pt.diag_test failed on device cuda
Labels: triaged,oncall: pt2

### 🐛 Describe the bug

The operator benchmark `pt.diag_test` fails on CUDA devices because the test argument [out_tensor](https://github.com/pytorch/pytorch/blob/3e48d927332915e1ecbd3c7f2c6b9680428f181e/benchmarks/operator_benchmark/pt/diag_test.py#L31-L33) is always allocated on the CPU, regardless of the specified device in the test configurations. This discrepancy leads to a device mismatch error when attempting to perform the operation on a CUDA device.

### Error logs

```
> python -m pt.diag_test
# ----------------------------------------
# PyTorch/Caffe2 Operator Micro-benchmarks
# ----------------------------------------
# Tag : short

# Benchmarking PyTorch: diag
# Mode: Eager
# Name: diag_dim1_M64_N64_diagonal0_outTrue_cpu
# Input: dim: 1, M: 64, N: 64, diagonal: 0, out: True, device: cpu
Forward Execution Time (us) : 4.015

# Benchmarking PyTorch: diag
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/home/tursom/dev/pytorch_raw/benchmarks/operator_benchmark/pt/diag_test.py", line 48, in <module>
    op_bench.benchmark_runner.main()
  File "/home/tursom/dev/pytorch_raw/benchmarks/operator_benchmark/benchmark_runner.py", line 166, in main
    benchmark_core.BenchmarkRunner(args).run()
  File "/home/tursom/dev/pytorch_raw/benchmarks/operator_benchmark/benchmark_core.py", line 436, in run
    launch_func(
  File "/home/tursom/dev/pytorch_raw/benchmarks/operator_benchmark/benchmark_core.py", line 275, in _launch_forward
    forward_time = timeit.timeit(
                   ^^^^^^^^^^^^^^
  File "/home/tursom/dev/pytorch_raw/.conda/lib/python3.12/timeit.py", line 237, in timeit
    return Timer(stmt, setup, timer, globals).timeit(number)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tursom/dev/pytorch_raw/.conda/lib/python3.12/timeit.py", line 180, in timeit
    timing = self.inner(it, self.timer)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<timeit-src>", line 6, in inner
  File "/home/tursom/dev/pytorch_raw/benchmarks/operator_benchmark/benchmark_pytorch.py", line 162, in run_forward
    self.output = self.op_bench.forward_impl()
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tursom/dev/pytorch_raw/benchmarks/operator_benchmark/benchmark_pytorch.py", line 68, in forward_impl
    return self.forward(*self.get_inputs())
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tursom/dev/pytorch_raw/benchmarks/operator_benchmark/pt/diag_test.py", line 39, in forward
    return torch.diag(input, diagonal=diagonal, out=out_tensor)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected out tensor to have device cuda:0, but got cpu instead
```

### Minified repro

```bash
cd benchmarks/operator_benchmark
python -m pt.diag_test
```

### Versions

Collecting environment information...
PyTorch version: 2.5.0a0+gitf7eae27
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04 LTS (x86_64)
GCC version: (Anaconda gcc) 11.2.0
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.39

Python version: 3.12.3 | packaged by Anaconda, Inc. | (main, May  6 2024, 19:46:43) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2060 SUPER
Nvidia driver version: 552.44
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             12
On-line CPU(s) list:                0-11
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen 5 3600 6-Core Processor
CPU family:                         23
Model:                              113
Thread(s) per core:                 2
Core(s) per socket:                 6
Socket(s):                          1
Stepping:                           0
BogoMIPS:                           8400.01
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 tsc_reliable nonstop_tsc cpuid extd_apicid 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 osvw topoext ssbd ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr virt_ssbd arat umip rdpid
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          192 KiB (6 instances)
L1i cache:                          192 KiB (6 instances)
L2 cache:                           3 MiB (6 instances)
L3 cache:                           16 MiB (1 instance)
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow: Mitigation; safe RET
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, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] mypy==1.9.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] optree==0.11.0
[pip3] torch==2.4.0a0+gita150df6
[conda] cudatoolkit               11.8.0               h6a678d5_0
[conda] magma-cuda118             2.6.1                         1    pytorch
[conda] mkl-include               2024.2.0              intel_663    intel
[conda] mkl-static                2024.2.0              intel_663    intel
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] torch                     2.4.0a0+gita150df6           dev_0    <develop>

cc @ezyang @anijain2305 @chauhang @penguinwu