MetricsCollection is not Pytorch JIT-able due to its metric_state property

[ethantang-db]

🐛 Bug

When attempting to run torch.jit.script on any modules that have a MetricsCollection object. You would run into an error something akin to the following

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/ethan.tang/source/composer/venv/lib/python3.9/site-packages/torch/jit/_script.py", line 1429, in script
    ret = _script_impl(
  File "/Users/ethan.tang/source/composer/venv/lib/python3.9/site-packages/torch/jit/_script.py", line 1147, in _script_impl
    return torch.jit._recursive.create_script_module(
  File "/Users/ethan.tang/source/composer/venv/lib/python3.9/site-packages/torch/jit/_recursive.py", line 557, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/Users/ethan.tang/source/composer/venv/lib/python3.9/site-packages/torch/jit/_recursive.py", line 634, in create_script_module_impl
    create_methods_and_properties_from_stubs(
  File "/Users/ethan.tang/source/composer/venv/lib/python3.9/site-packages/torch/jit/_recursive.py", line 466, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(
RuntimeError: 
items method does not accept any arguments:
  File "/Users/ethan.tang/source/composer/venv/lib/python3.9/site-packages/torchmetrics/collections.py", line 193
    def metric_state(self) -> Dict[str, Dict[str, Any]]:
        """Get the current state of the metric."""
        return {k: m.metric_state for k, m in self.items(keep_base=False, copy_state=False)}
                                              ~~~~~~~~~~ <--- HERE

To Reproduce

Steps to reproduce the behavior...

Code sample

from torchmetrics import MetricCollection
from torchmetrics.classification import MulticlassAccuracy
import torch

metrics = MulticlassAccuracy(num_classes=10, average='micro') 
metric_collection = MetricCollection([metrics]) 
exported_collection = torch.jit.script(metric_collection)  # error happens here

Expected behavior

This should run and compile properly without error.

Environment

• TorchMetrics version (if build from source, add commit SHA): 1.4.1 -> 1.5.1
• Python & PyTorch Version (e.g., 1.0): 2.5.0
• Any other relevant information such as OS (e.g., Linux): Mac OS X Sonoma, although also reproduced on Linux based CI machines

Additional context

Most likely the issue is due to

    @property
    def metric_state(self) -> Dict[str, Dict[str, Any]]:
        """Get the current state of the metric."""
        return {k: m.metric_state for k, m in self.items(keep_base=False, copy_state=False)}

When I comment out @property, this works but it would also break the style of acquiring metric_state like it is with Metrics class. It can also be due to the fact that the item() property has been redefined and torch JIT does not like that. Not sure what the best fix is for this though.

[github-actions]

Hi! thanks for your contribution!, great first issue!