Repository: pytorch
Issue: ONNX export: w=tensor.shape() in torch.arange(end=w...) causes "tensor does not have a device" error
Labels: module: onnx,triaged

### 🐛 Describe the bug

The error looks like this:

```
Traceback (most recent call last):
  File "/mnt/dingus_drive/catid/train_detector/onnx_export.py", line 119, in export_to_onnx
    torch.onnx.export(
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/utils.py", line 551, in export
    _export(
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/utils.py", line 1648, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/utils.py", line 1174, in _model_to_graph
    graph = _optimize_graph(
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/utils.py", line 714, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/utils.py", line 1997, in _run_symbolic_function
    return symbolic_fn(graph_context, *inputs, **attrs)
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/symbolic_opset11.py", line 864, in arange
    return g.op("Range", start_default, end, delta_default)
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/_internal/jit_utils.py", line 93, in op
    return _add_op(self, opname, *raw_args, outputs=outputs, **kwargs)
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/_internal/jit_utils.py", line 252, in _add_op
    node = _create_node(
  File "/home/saronic/miniconda3/envs/train/lib/python3.10/site-packages/torch/onnx/_internal/jit_utils.py", line 314, in _create_node
    _C._jit_pass_onnx_node_shape_type_inference(node, params_dict, opset_version)
RuntimeError: tensor does not have a device
```

Seems like the ONNX export C++ code does not have support for this particular thing:

If you take a tensor.shape, and use the result as the end= parameter for torch.arange() it leads to this internal error.  I'm working around it like this, explicitly creating a tensor object from the shape:

```
shapes = torch.tensor([(f.shape[2], f.shape[3]) for f in thing], device=device)
```

### Versions

Collecting environment information...
PyTorch version: 2.4.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.13 | packaged by conda-forge | (main, Dec 23 2023, 15:36:39) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-169-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.6.20
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100 80GB PCIe
GPU 1: NVIDIA A100 80GB PCIe
GPU 2: NVIDIA A100 80GB PCIe
GPU 3: NVIDIA A100 80GB PCIe
GPU 4: NVIDIA A100 80GB PCIe
GPU 5: NVIDIA A100 80GB PCIe
GPU 6: NVIDIA A100 80GB PCIe
GPU 7: NVIDIA A100 80GB PCIe

Nvidia driver version: 560.28.03
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, 57 bits virtual
CPU(s):                             96
On-line CPU(s) list:                0-95
Thread(s) per core:                 2
Core(s) per socket:                 6
Socket(s):                          8
NUMA node(s):                       2
Vendor ID:                          GenuineIntel
CPU family:                         6
Model:                              106
Model name:                         Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz
Stepping:                           6
CPU MHz:                            1999.999
BogoMIPS:                           3999.99
Virtualization:                     VT-x
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          1.5 MiB
L1i cache:                          1.5 MiB
L2 cache:                           192 MiB
L3 cache:                           128 MiB
NUMA node0 CPU(s):                  0-47
NUMA node1 CPU(s):                  48-95
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        Not affected
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:             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; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
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 tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves wbnoinvd arat avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid md_clear arch_capabilities

Versions of relevant libraries:
[pip3] adam-atan2-pytorch==0.0.12
[pip3] numpy==1.26.4
[pip3] onnx==1.16.2
[pip3] onnxoptimizer==0.3.13
[pip3] onnxruntime-gpu==1.19.2
[pip3] onnxslim==0.1.34
[pip3] pytorch-ranger==0.1.1
[pip3] torch==2.4.1
[pip3] torch-optimizer==0.3.1a0
[pip3] torchvision==0.19.1
[pip3] triton==3.0.0
[conda] adam-atan2-pytorch        0.0.12                   pypi_0    pypi
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] pytorch-ranger            0.1.1                    pypi_0    pypi
[conda] torch                     2.4.1                    pypi_0    pypi
[conda] torch-optimizer           0.3.1a0                  pypi_0    pypi
[conda] torchvision               0.19.1                   pypi_0    pypi
[conda] triton                    3.0.0                    pypi_0    pypi