Fix moving keys to device in ResultCollection

[clumsy]

What does this PR do?

Move registered results' keys to the current device. Otherwise internal tensors like cumulated_batch_size will stay on cpu after they are moved there in previous call to teardown().

Removing standalone register_key() method because it's not used outside of this module and we either end up calling .to() twice or might need to require the potential users to call it themselves.

Fixes #19813

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

---

📚 Documentation preview 📚: https://pytorch-lightning--19814.org.readthedocs.build/en/19814/

[clumsy]

Please check if you have a chance, @awaelchli

[clumsy]

Could you please check this bug, @Borda?

[clumsy]

I would appreciate if you can check this PR, @awaelchli. Thanks!

[awaelchli]

@clumsy with these small adjustments I commented on, this would be ready to land. I can also make the changes if you don't have time. Anyway, sorry for the delay.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81%. Comparing base (2064887) to head (0bd1563).

❗ There is a different number of reports uploaded between BASE (2064887) and HEAD (0bd1563). Click for more details.

HEAD has 274 uploads less than BASE

Flag	BASE (2064887)	HEAD (0bd1563)
cpu	83	21
lightning	61	16
pytest	46	2
python3.9	23	6
python3.10	24	6
lightning_fabric	13	0
python3.11	24	6
gpu	4	2
python3.12	12	3
pytorch2.1	23	12
pytest-full	41	21
pytorch2.2	6	3
pytorch2.3	6	3
pytorch_lightning	13	7
pytorch2.4	6	3

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #19814     +/-   ##
=========================================
- Coverage      89%      81%     -8%     
=========================================
  Files         266      263      -3     
  Lines       22976    22920     -56     
=========================================
- Hits        20522    18556   -1966     
- Misses       2454     4364   +1910     

[awaelchli]

Great fix, thanks @clumsy!