[5676209][ONNX][Autocast] Add check for input bs vs calibration data bs

[gcunhase]

What does this PR do?

Type of change: Bug fix

Overview: Autocast crashes if the input batch size in the ONNX model is different to the calibration data input batch size. For example: calibration data has shape [10, 6, 3, 480, 800] and ONNX model has shape [1, 6, 3, 480, 800]. The quantization workflow interprets this as 10 calibration samples, so ideally, Autocast would also interpret them similarly. This PR just allows Autocast to exit gracefully with a custom message.

Usage

$ python -m modelopt.onnx.autocast --onnx_path=$MODEL_NAME.onnx --calibration_data=calib_data_10.npz

Testing

See bug 5676209.

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: No
• Did you add or update any necessary documentation?: Yes
• Did you update Changelog?: No

Additional Information

Original error:

polygraphy.exception.exception.PolygraphyException: Input tensor: image | Received incompatible shape: (10, 6, 3, 480, 800).
Note: Expected a shape compatible with: BoundedShape([1, 6, 3, 480, 800], min=None, max=None)

Autocast error:

ValueError: Input shape from 'image' does not match provided input shape: [1, 6, 3, 480, 800] vs [10, 6, 3, 480, 800]. Please make sure that your calibration data matches the ONNX input shapes.

[galagam]

Thanks for this fix!