Repository: pytorch
Issue: [FakeTensor] Error: SymIntArrayRef expected to contain only concrete integers in UpSample2D block
Labels: module: onnx,oncall: export

### 🐛 Describe the bug

I was trying to export the DiffSTE model from this [repo](https://github.com/UCSB-NLP-Chang/DiffSTE) to onnx using **dynamo_export**. The model is able to be trained and run inference normally. However, when I export the model, it complains that `RuntimeError: aten/src/ATen/RegisterCompositeExplicitAutograd.cpp:2305: SymIntArrayRef expected to contain only concrete integers` for the UpBlock2D, which is calling the Upsample2D subsequently. I also attach the Dynmo Export report below.
```python
        # diffusers/models/unets/unet_2d_blocks.py", line 2685
        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
                hidden_states = upsampler(hidden_states, upsample_size)        <--------------- this part

        return hidden_states
```
```python
       # diffusers/models/upsampling.py", line 171, in forward
       hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest") <--------------- this part
```

# Export Code
``` python

def onnx_export(
    model,
    model_args: tuple,
    output_path: Path,
    ordered_input_names,
    output_names,
    dynamic_axes,
    opset,
    use_external_data_format=False,
):
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with torch.no_grad():
        model.eval()
        export_options = torch.onnx.ExportOptions(dynamic_shapes=True)
        if isinstance(model_args, tuple):
            torch.onnx.dynamo_export(
                model, 
                *model_args,
                export_options=export_options
            ).save(output_path.as_posix())
        else:
            torch.onnx.dynamo_export(
                model, 
                model_args,
                export_options=export_options
            ).save(output_path.as_posix())

# UNET
unet_in_channels = model.unet.config.in_channels
unet_path = output_path / "unet" / "model.onnx"
onnx_export(
    model.unet,
    model_args=(
        torch.randn(2, unet_in_channels, img_size // 8, img_size // 8).to(device=device, dtype=dtype),
        torch.tensor(100, dtype=torch.float32),
        torch.randn(2, num_tokens, config.model.unet.cross_attention_dim["text"]).to(device=device, dtype=dtype),
        torch.randn(2, num_tokens, config.model.unet.cross_attention_dim["char"]).to(device=device, dtype=dtype),  
    ),
    output_path=unet_path,
    ordered_input_names=["sample", "timestep", "text_hidden_states", "char_hidden_states"],
    output_names=["out_sample"],  # has to be different from "sample" for correct tracing
    dynamic_axes={
        "sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
        "timestep": {},
        "text_hidden_states": {0: "batch", 1: "sequence"},
        "char_hidden_states": {0: "batch", 1: "sequence"},
    },
    opset=opset,
    use_external_data_format=True,  # UNet is > 2GB, so the weights need to be split
)
unet_model_path = str(unet_path.absolute().as_posix())
unet_dir = os.path.dirname(unet_model_path)
unet = onnx.load(unet_model_path)
# clean up existing tensor files
shutil.rmtree(unet_dir)
os.mkdir(unet_dir)
# collate external tensor files into one
onnx.save_model(
    unet,
    unet_model_path,
    save_as_external_data=True,
    all_tensors_to_one_file=True,
    location="weights.pb",
    convert_attribute=False,
)
del model.unet
```


# Some stack traces
```
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761] failed while attempting to run meta for aten._unsafe_index.Tensor
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761] Traceback (most recent call last):
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]   File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py", line 1757, in _dispatch_impl
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]     r = func(*args, **kwargs)
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]   File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_ops.py", line 667, in __call__
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]     return self_._op(*args, **kwargs)
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]   File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_meta_registrations.py", line 3019, in meta_index_Tensor
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]     indices = list(refs._maybe_broadcast(*indices))
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]   File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_refs/__init__.py", line 439, in _maybe_broadcast
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]     return tuple(__maybe_broadcast(x, common_shape) for x in args)
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]   File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_refs/__init__.py", line 439, in <genexpr>
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]     return tuple(__maybe_broadcast(x, common_shape) for x in args)
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]   File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_refs/__init__.py", line 431, in __maybe_broadcast
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761]     return x.expand(common_shape)
E0822 23:10:34.280457 139634357630784 torch/_subclasses/fake_tensor.py:1761] RuntimeError: aten/src/ATen/RegisterCompositeExplicitAutograd.cpp:2305: SymIntArrayRef expected to contain only concrete integers
Traceback (most recent call last):
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 1504, in dynamo_export
    return Exporter(
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 1236, in export
    graph_module = self.options.fx_tracer.generate_fx(
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/fx/dynamo_graph_extractor.py", line 233, in generate_fx
    return self.pre_export_passes(options, model, graph_module, updated_model_args)  # type: ignore[return-value]
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/fx/dynamo_graph_extractor.py", line 243, in pre_export_passes
    return exporter.common_pre_export_passes(
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 1548, in common_pre_export_passes
    module = passes.Functionalize(
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 152, in wrapper
    ctx.log_and_raise_if_error(diag)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 369, in log_and_raise_if_error
    raise diagnostic.source_exception
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 136, in wrapper
    return_values = fn(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/fx/_pass.py", line 278, in run
    module = self._run(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/fx/passes/functionalization.py", line 124, in _run
    graph_module = proxy_tensor.make_fx(
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 1421, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 1367, in trace
    return self._trace_inner(f, *args)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 1354, in _trace_inner
    t = dispatch_trace(
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_compile.py", line 31, in inner
    return disable_fn(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 600, in _fn
    return fn(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 642, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 600, in _fn
    return fn(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/_symbolic_trace.py", line 822, in trace
    (self.create_arg(fn(*args)),),
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 660, in wrapped
    out = f(*tensors)
  File "<string>", line 1, in <lambda>
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/fx/passes/functionalization.py", line 87, in wrapped
    out = function(*inputs_functional)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/onnx/_internal/fx/passes/_utils.py", line 31, in wrapped
    return torch.fx.Interpreter(graph_module).run(*args)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/interpreter.py", line 146, in run
    self.env[node] = self.run_node(node)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/interpreter.py", line 203, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/interpreter.py", line 275, in call_function
    return target(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_ops.py", line 667, in __call__
    return self_._op(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 705, in __torch_function__
    return func(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_ops.py", line 667, in __call__
    return self_._op(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/utils/_stats.py", line 21, in wrapper
    return fn(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 755, in __torch_dispatch__
    return self.inner_torch_dispatch(func, types, args, kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 790, in inner_torch_dispatch
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py", line 467, in proxy_call
    out = func(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_ops.py", line 667, in __call__
    return self_._op(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/utils/_stats.py", line 21, in wrapper
    return fn(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py", line 1061, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py", line 1450, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py", line 1153, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_subclasses/fake_tensor.py", line 1757, in _dispatch_impl
    r = func(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_ops.py", line 667, in __call__
    return self_._op(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_meta_registrations.py", line 3019, in meta_index_Tensor
    indices = list(refs._maybe_broadcast(*indices))
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_refs/__init__.py", line 439, in _maybe_broadcast
    return tuple(__maybe_broadcast(x, common_shape) for x in args)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_refs/__init__.py", line 439, in <genexpr>
    return tuple(__maybe_broadcast(x, common_shape) for x in args)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/_refs/__init__.py", line 431, in __maybe_broadcast
    return x.expand(common_shape)
RuntimeError: aten/src/ATen/RegisterCompositeExplicitAutograd.cpp:2305: SymIntArrayRef expected to contain only concrete integers

While executing %_unsafe_index : [num_users=1] = call_function[target=torch.ops.aten._unsafe_index.Tensor](args = (%div_41, [None, None, %unsqueeze_29, %_to_copy_2]), kwargs = {})
Original traceback:
  File "/home/ltruong/synthetic_license_plate/models/diffste/unet_2d_multicondition.py", line 387, in forward
    sample = upsample_block(
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/diffusers/models/unets/unet_2d_blocks.py", line 2685, in forward
    hidden_states = upsampler(hidden_states, upsample_size)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ltruong/miniconda3/envs/synth/lib/python3.8/site-packages/diffusers/models/upsampling.py", line 171, in forward
    hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
```

cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @penguinwu @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @eellison



### Versions

```
[pip3] numpy==1.24.4
[pip3] onnx==1.15.0
[pip3] onnx-graphsurgeon==0.5.2
[pip3] onnxoptimizer==0.3.13
[pip3] onnxruntime==1.19.0
[pip3] onnxscript==0.1.0.dev20240822
[pip3] pytorch-lightning==2.4.0
[pip3] torch==2.4.0+cu124
[pip3] torch-fidelity==0.3.0
[pip3] torchaudio==2.4.0+cu124
[pip3] torchmetrics==1.4.0.post0
[pip3] torchtext==0.18.0
[pip3] torchvision==0.19.0+cu124
[pip3] triton==3.0.0
[pip3] tritonclient==2.48.0
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] pytorch-lightning         2.4.0                    pypi_0    pypi
[conda] torch                     2.4.0+cu124              pypi_0    pypi
[conda] torch-fidelity            0.3.0                    pypi_0    pypi
[conda] torchaudio                2.4.0+cu124              pypi_0    pypi
[conda] torchmetrics              1.4.0.post0              pypi_0    pypi
[conda] torchtext                 0.18.0                   pypi_0    pypi
[conda] torchvision               0.19.0+cu124             pypi_0    pypi
[conda] triton                    3.0.0                    pypi_0    pypi
[conda] tritonclient              2.48.0                   pypi_0    pypi
```