[Feat] Add support for distributed training

[tchaton]

Description

This PR adds support for computing support on new data for running inference. This is only supported by LightningPrototypicalNetworks right now, as the classifier needs re-fitting.

This is used for Lightning Flash learn2learn integration: https://github.com/PyTorchLightning/lightning-flash/pull/737/

Replace this line by a one sentence summary of your PR.

If necessary, use the following space to provide context or more details.

Contribution Checklist

If your contribution modifies code in the core library (not docs, tests, or examples), please fill the following checklist.

• [x] My contribution is listed in CHANGELOG.md with attribution.
• [x] My contribution modifies code in the main library.
• [ ] My modifications are tested.
• [ ] My modifications are documented.

Optional

If you make major changes to the core library, please run make alltests and copy-paste the content of alltests.txt below.

[PASTE HERE]