fix: ND input quantization

[narendasan]

Description

Flattens ND quantization inputs to 2D

Fixes #4201

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/models/test_models_export.py	2026-04-23 17:31:15.121758+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/models/test_models_export.py	2026-04-23 17:31:41.297024+00:00
@@ -409,13 +409,13 @@
                use_explicit_typing=True,
            )
            expected = model(input_tensor)
            outputs_trt = trt_model(input_tensor)
            abs_diff = torch.abs(expected - outputs_trt)
-            assert torch.allclose(expected, outputs_trt, rtol=0.3, atol=0.3), (
-                f"max abs diff: {abs_diff.max().item()}"
-            )
+            assert torch.allclose(
+                expected, outputs_trt, rtol=0.3, atol=0.3
+            ), f"max abs diff: {abs_diff.max().item()}"


@unittest.skipIf(
    torch.cuda.get_device_capability() < (8, 9),
    "FP8 quantization requires compute capability 8.9 or later",
@@ -472,13 +472,13 @@
                use_explicit_typing=True,
            )
            expected = model(input_tensor)
            outputs_trt = trt_model(input_tensor)
            abs_diff = torch.abs(expected - outputs_trt)
-            assert torch.allclose(expected, outputs_trt, rtol=5e-2, atol=5e-2), (
-                f"max abs diff: {abs_diff.max().item()}"
-            )
+            assert torch.allclose(
+                expected, outputs_trt, rtol=5e-2, atol=5e-2
+            ), f"max abs diff: {abs_diff.max().item()}"


@unittest.skipIf(
    torch.cuda.get_device_capability() < (8, 9),
    "FP8 quantization requires compute capability 8.9 or later",

[zewenli98]

LGTM