fix: prob attention shape error while bs>1

[LongxingTan]

Description

This PR ...

Checklist

• Linked issues (if existing)
• Amended changelog.md for large changes (and added myself there as contributor)
• Added/modified tests
• Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
  To run hooks independent of commit, execute pre-commit run --all-files

Thank you for joining. Have fun coding!

[codecov]

Codecov Report

All modified lines are covered by tests ✅

Comparison is base (679e349) 89.07% compared to head (57095bd) 89.07%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master      #50      +/-   ##
==========================================
- Coverage   89.07%   89.07%   -0.01%     
==========================================
  Files          27       27              
  Lines        2133     2132       -1     
==========================================
- Hits         1900     1899       -1     
  Misses        233      233              

Flag	Coverage Δ
cpu	89.07% <100.00%> (-0.01%)	⬇️
unittest	89.07% <100.00%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files	Coverage Δ
tfts/layers/attention_layer.py	89.93% <100.00%> (-0.07%)	⬇️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.