fix: Reduce GPU to CPU copies at inference

[sjrl]

Related Issues

• fixes Copy less from GPU to memory when postprocessing QA logits #2926

Proposed Changes:

Implemented the recommended fix described in the issue. Basically, we perform one GPU to CPU copy of entire matrices instead of individually copying over elements of the matrix.

How did you test it?

Notes for the reviewer

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added tests that demonstrate the correct behavior of the change
• I've used the conventional commit convention for my PR title
• I documented my code
• I ran pre-commit hooks and fixed any issue

[sjrl]

Comment: #2926 (comment) provides timings for how these changes improve the speed of the reader predictions when using large values ~50 for top_k_per_sample.

[sjrl]

Hi, @bogdankostic I want to let you know that this PR is ready for review. The comment referenced above shows how the changes improve the speed of the reader when using large (e.g. 50) values for top_k_per_sample. There are no more obvious speedups that I can see and if we need to we can look at further optimizing this function in the future.

[bogdankostic]

Wow, these speed improvements seem impressive 🚀 LGTM!