Fix Qwen3 ONNX export bug related to 0/1 specialization

[victoroliv2]

What does this PR do?

**Type of change: ** Bug Fix

In PyTorch 2.7 (version I am using), Dynamic batching in ONNX export is broken with dimension 0/1. I believe this has been fixed in pytorch/pytorch#184837

This PR fixes this issue by using dim = 2 so Qwen3 ONNX export works on older PyTorch releases (my use case).

Overview: ?

Usage

# Add a code snippet demonstrating how to use this

🚀 Pull Request Checklist

Thank you for contributing to TensorRT Edge-LLM! Before we review your pull request, please make sure the following items are complete.
Please also refer to Contributor guidelines for general guidelines.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit.
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing.

📄 Documentation

• Updated any necessary documentation

⚙️ Compatibility

• The change is backward compatible

Additional Information