Update assertion and argument description

Signed-off-by: ajrasane <[email protected]>

Author: ajrasane

examples/onnx_ptq/torch_quant_to_onnx.py

@@ -123,7 +123,7 @@ def main():
         "--batch_size",
         type=int,
         default=1,
-        help="Batch size for calibration.",
+        help="Batch size for calibration and ONNX model export.",
     )
 
     args = parser.parse_args()

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -486,7 +486,7 @@ def get_onnx_bytes_and_metadata(
         param_dtype = torch.float32
     if weights_dtype in ["fp16", "bf16"] and param_dtype == torch.float32:
         if is_mxfp8_quantized(model) or is_int4_quantized(model):
-            assert weights_dtype == "fp16", "BF16 + MXFP8 mixed precision is not supported yet"
+            assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet"
             onnx_opt_graph = convert_float_to_float16(
                 onnx_opt_graph,
                 keep_io_types=False,