Repository: pytorch
Issue: ONNX NMS Operator Output Different From Torchvision Implementation
Labels: module: onnx,triaged

### 🐛 Describe the bug

I am trying to encapsulate the torchvision.ops.nms function in an onnx model, there is no problem in the model conversion and inference but the output of the derived onnx model differs from the torchvision implementation.
```py
import torch
import torchvision
import onnxruntime
import numpy as np

class NMS(torch.nn.Module):
    def __init__(self, iou_threshold=0.45):
        super(NMS, self).__init__()
        self.iou_threshold = iou_threshold

    def forward(self, x):
        boxes = x[:, :4]
        scores = x[:, 4]
        keep = torchvision.ops.nms(boxes, scores, self.iou_threshold)
        return keep

# Test data
nms_x = torch.rand(50, 38)

# PyTorch model
nms_model = NMS()
torch_output = nms_model(nms_x)

# Export to ONNX
torch.onnx.export(nms_model, 
                  (nms_x,),
                  "nms.onnx",
                  opset_version=17, 
                  input_names=["input"],
                  output_names=["output"],
                  #dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
                  )

# ONNX Runtime inference
session = onnxruntime.InferenceSession("nms.onnx")

onnx_output = session.run(None, {'input': nms_x.numpy()})[0]

print("PyTorch output shape:", torch_output.shape)
print("ONNX output shape:", onnx_output.shape)
```
![image](https://github.com/user-attachments/assets/594bb611-e250-40b2-93c3-3ecb2aa10046)


### Versions

Collecting environment information...
PyTorch version: 2.4.0+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Home
GCC version: (tdm64-1) 10.3.0
Clang version: 17.0.6
CMake version: version 3.30.1
Libc version: N/A

Python version: 3.11.5 | packaged by Anaconda, Inc. | (main, Sep 11 2023, 13:26:23) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22631-SP0
Is CUDA available: False
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4070 Laptop GPU
Nvidia driver version: 560.94
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=4001
DeviceID=CPU0
Family=107
L2CacheSize=8192
L2CacheSpeed=
Manufacturer=AuthenticAMD
MaxClockSpeed=4001
Name=AMD Ryzen 9 7940HS w/ Radeon 780M Graphics
ProcessorType=3
Revision=29697

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] numpydoc==1.5.0
[pip3] onnx==1.16.2
[pip3] onnx-graphsurgeon==0.3.27
[pip3] onnx-simplifier==0.4.36
[pip3] onnx-tool==0.9.0
[pip3] onnx2pytorch==0.4.1
[pip3] onnx2tf==1.19.11
[pip3] onnxoptimizer==0.3.13
[pip3] onnxruntime==1.18.1
[pip3] onnxscript==0.1.0.dev20240823
[pip3] onnxsim==0.4.35
[pip3] sclblonnx==0.2.1
[pip3] sng4onnx==1.0.1
[pip3] tf2onnx==1.16.1
[pip3] torch==2.4.0
[pip3] torchvision==0.19.0
[conda] _anaconda_depends 2023.09 py311_mkl_1
[conda] blas 1.0 mkl
[conda] mkl 2021.4.0 pypi_0 pypi
[conda] mkl-service 2.4.0 py311h2bbff1b_1
[conda] mkl_fft 1.3.8 py311h2bbff1b_0
[conda] mkl_random 1.2.4 py311h59b6b97_0
[conda] numpy 1.26.4 pypi_0 pypi
[conda] numpydoc 1.5.0 py311haa95532_0
[conda] onnx2pytorch 0.4.1 pypi_0 pypi
[conda] torch 2.4.0 pypi_0 pypi
[conda] torchvision 0.19.0 pypi_0 pypi