Fix support for ET repo generated pte by optionally adding batch dim

[Jack-Khuu]

As titled, torchchat expects logits as (1, 1, vocab size), but depending on the exporting logic the batch dim is not provided (for example with pte generated from the ET repo

This just adds a simple check and layer if needed.

Tested by

python3 torchchat.py generate llama3.1 --pte-path <pte from ET repo export>
python3 torchchat.py generate llama3.1 --pte-path <pte from TC repo export>

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1177

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 0d700de with merge base 971ed93 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.