Repository: pytorch
Issue: Type parameter (T) of Optype (Where) bound to different types (tensor(int64) and tensor(float) in node (/Where)
Labels: module: onnx,triaged,OSS contribution wanted

### 🐛 Describe the bug

I'm getting an invalid onnx model (full PyTorch code for generating it at the bottom). When I attempt to run it with polygraphy (`polygraphy run traversability.onnx --onnxrt`), I get

```
  File "/home/dspyz/.local/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 419, in __init__
    self._create_inference_session(providers, provider_options, disabled_optimizers)
  File "/home/dspyz/.local/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 472, in _create_inference_session
    sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from /home/dspyz/foo/traversability.onnx failed:Type Error: Type parameter (T) of Optype (Where) bound to different types (tensor(int64) and tensor(float) in node (/Where).
```

[traversability.zip](https://github.com/pytorch/pytorch/files/14873836/traversability.zip)

```
import torch
import torch.nn as nn
import torch.onnx

MAX_TRAV = 21


class Traversability(nn.Module):
    def forward(self, untraversable):
        max_trav_scalar = torch.tensor(
            MAX_TRAV, dtype=torch.float32, device=untraversable.device
        )
        sq_distances = torch.where(untraversable, 0, max_trav_scalar**2)
        extra_column = torch.full(
            (sq_distances.shape[0], 1),
            max_trav_scalar**2,
            device=sq_distances.device,
        )
        # horizontal pass
        for d in range(1, MAX_TRAV * 2 + 1, 2):
            d = torch.tensor(d, dtype=torch.int32, device=sq_distances.device)
            sq_distances = torch.minimum(
                sq_distances,
                torch.minimum(
                    torch.cat([sq_distances[:, 1:] + d, extra_column], 1),
                    torch.cat([extra_column, sq_distances[:, :-1] + d], 1),
                ),
            )
        extra_row = torch.full(
            (1, sq_distances.shape[1]),
            max_trav_scalar**2,
            device=sq_distances.device,
        )
        # vertical pass
        for d in range(1, MAX_TRAV * 2 + 1, 2):
            d = torch.tensor(d, dtype=torch.int32, device=sq_distances.device)
            sq_distances = torch.minimum(
                sq_distances,
                torch.minimum(
                    torch.cat([sq_distances[1:, :] + d, extra_row], 0),
                    torch.cat([extra_row, sq_distances[:-1, :] + d], 0),
                ),
            )
        return sq_distances.sqrt()


# Create an instance of the model
model = Traversability()


def save_onnx():
    # Set the model to evaluation mode
    model.eval()

    # Define the input shape
    dummy_input = torch.randint(0, 2, (224, 224), dtype=torch.bool)

    # Specify the output file path
    onnx_file_path = "traversability.onnx"

    # Convert the model to ONNX format
    torch.onnx.export(
        model,
        dummy_input,
        onnx_file_path,
        input_names=["input"],
        output_names=["output"],
    )

    print("Model converted to ONNX format. Saved as:", onnx_file_path)


if __name__ == "__main__":
    save_onnx()
```


### Versions

Versions of relevant libraries:
[pip3] mypy==1.3.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.2
[pip3] onnx==1.15.0
[pip3] onnxruntime==1.17.1
[pip3] torch==2.2.1
[pip3] triton==2.2.0
[conda] Could not collect