export.py: fix custom SDPA type conversion logic & re-enable for bfloat16

[swolchok]

q.dtype is always torch.float at this point. Clearly meant to use the input dtype.

[swolchok]

Stack from ghstack (oldest at bottom):

• Add bf16 to runner_et test #1196
• -> export.py: fix custom SDPA type conversion logic & re-enable for bfloat16 #1193

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 6a1ff58 with merge base 6d2ef4a ():

NEW FAILURES - The following jobs have failed:

• pull / test-mps / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• pull / test-mps-dtype / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

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