multi-gpu inference. Adds 'batch index' to the resulting prediction

[skothenhill-nv]

Description

Adds 'batch index' to the resulting prediction dictionary. this allows users to reconstruct the original ordering of predictions with multi-gpu inference. Prevents users from using 'epoch' mode of inference with multiple GPUs. this addresses a known issue: https://nvbugswb.nvidia.com/NvBugs5/SWBug.aspx?bugid=4717442&cmtNo=

Type of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor
• Documentation update
• Other (please describe):

Pre-submit Checklist

• [x ] I have tested these changes locally
• I have updated the documentation accordingly
• All existing tests pass successfully

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 83.33333% with 1 line in your changes missing coverage. Please review.

Project coverage is 84.40%. Comparing base (29ce3bc) to head (924d426).

✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
...ges/bionemo-llm/src/bionemo/llm/utils/callbacks.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #854      +/-   ##
==========================================
- Coverage   84.40%   84.40%   -0.01%     
==========================================
  Files         138      138              
  Lines        8685     8690       +5     
==========================================
+ Hits         7331     7335       +4     
- Misses       1354     1355       +1     

Files with missing lines	Coverage Δ
...ges/bionemo-llm/src/bionemo/llm/utils/callbacks.py	94.44% <83.33%> (-2.33%)	⬇️