Repository: pytorch
Issue: [ONNX][dynamo] Failed to export cumsum with dtype=float16
Labels: module: onnx,triaged,module: half

### 🐛 Describe the bug

This is exposed by the test case: `test_output_match` in test/onnx/test_fx_op_consistency.py. When export cumsum with inputs of dtype=torch.float16 to an ONNX graph, it will get the following error:
`onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Exception during initialization: /onnxruntime_src/onnxruntime/core/framework/allocation_planner.cc:230 int& onnxruntime::PlannerImpl::UseCount(onnxruntime::OrtValueIndex) n >= 0 && static_cast<size_t>(n) < ort_value_info_.size() was false. invalid value index: -1 against size 5`

A simple reproducer:
```
import torch
import onnxruntime
import io
import numpy as np
# import intel_extension_for_pytorch as ipex
from typing import (
    Any,
    Callable,
    Mapping,
    Optional,
    Sequence,
    Union,
)
from torch.types import Number
_NumericType = Union[Number, torch.Tensor, np.ndarray]
_InputArgsType = Optional[
    Union[torch.Tensor, int, float, bool, Sequence[Any], Mapping[str, Any]]
]
_OutputsType = Sequence[_NumericType]

def run_ort(
    onnx_model: Union[str, torch.onnx.ExportOutput],
    pytorch_inputs: Sequence[_InputArgsType],
) -> _OutputsType:
    """Run ORT on the given ONNX model and inputs

    Used in test_fx_to_onnx_with_onnxruntime.py

    Args:
        onnx_model (Union[str, torch.onnx.ExportOutput]): Converter ONNX model
        pytorch_inputs (Sequence[_InputArgsType]): The given torch inputs

    Raises:
        AssertionError: ONNX and PyTorch should have the same input sizes

    Returns:
        _OutputsType: ONNX model predictions
    """
    if isinstance(onnx_model, torch.onnx.ExportOutput):
        buffer = io.BytesIO()
        onnx_model.save(buffer)
        ort_model = buffer.getvalue()
    else:
        ort_model = onnx_model

    # Suppress floods of warnings from ONNX Runtime
    session_options = onnxruntime.SessionOptions()
    session_options.log_severity_level = 3  # Error
    session = onnxruntime.InferenceSession(
        ort_model, providers=["CPUExecutionProvider"], sess_options=session_options
    )
    input_names = [ort_input.name for ort_input in session.get_inputs()]

    if len(input_names) != len(pytorch_inputs):
        raise AssertionError(
            f"Expected {len(input_names)} inputs, got {len(pytorch_inputs)}"
        )

    ort_input = {k: v.cpu().numpy() for k, v in zip(input_names, pytorch_inputs)}
    return session.run(None, ort_input)

class SingleOpModel(torch.nn.Module):
    """Test model to wrap around a single op for export."""

    def __init__(self, op, kwargs):
        super().__init__()
        self.operator = op
        self.kwargs = kwargs

    def forward(self, *args):
        return self.operator(*args, **self.kwargs)


set_dtype = torch.float16

input = (torch.randn((5, 5, 5), dtype=set_dtype), 1)
kwargs = {"dtype": set_dtype}
model = SingleOpModel(torch.cumsum, kwargs)
ref_input_kwargs = {}

export_output = torch.onnx.dynamo_export(
    model,
    *input,
    **ref_input_kwargs,
    export_options=torch.onnx.ExportOptions(
        op_level_debug=False,
        dynamic_shapes=False,
        diagnostic_options=torch.onnx.DiagnosticOptions(
            verbosity_level=10
        ),
    ),
)
onnx_format_args = export_output.adapt_torch_inputs_to_onnx(
        *input, **ref_input_kwargs
    )
ort_outputs = run_ort(export_output, onnx_format_args)
```

### Versions

PyTorch version: 2.2.0a0+git5fa0c13
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 8.5.2111 (x86_64)
GCC version: (GCC) 11.2.1 20210728 (Red Hat 11.2.1-1)
Clang version: 16.0.0 (Red Hat 16.0.0-2.module_el8+405+25122a8c)
CMake version: version 3.21.4
Libc version: glibc-2.28

Python version: 3.8.5 (default, Sep  4 2020, 07:30:14)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.16.0-rc1-intel-next-00543-g5867b0a2a125-x86_64-with-glibc2.10
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] flake8==3.8.2
[pip3] flake8-bugbear==20.1.4
[pip3] flake8-coding==1.3.3
[pip3] flake8-comprehensions==3.3.0
[pip3] flake8-executable==2.0.4
[pip3] flake8-pyi==20.5.0
[pip3] intel-extension-for-pytorch==2.2.0+gite7090c6
[pip3] mypy==1.4.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.22.4
[pip3] onnx==1.14.1
[pip3] onnxruntime==1.15.1
[pip3] onnxscript==0.1.0.dev20230830
[pip3] torch==2.2.0a0+git46b6478
[pip3] torchvision==0.16.0a0+fb115c2
[pip3] triton==2.0.0
[conda] intel-extension-for-pytorch 2.2.0+gite7090c6           dev_0    <develop>
[conda] mkl                       2022.1.0           hc2b9512_224  
[conda] mkl-include               2023.1.0                 pypi_0    pypi
[conda] mkl-static                2023.1.0                 pypi_0    pypi
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] torch                     2.2.0a0+git46b6478           dev_0    <develop>
[conda] torchvision               0.16.0a0+fb115c2           dev_0    <develop>
[conda] triton                    2.0.0                    pypi_0    pypi