Fix getitem for metric collection when prefix/postfix is set

CHANGELOG.md

@@ -36,6 +36,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed dtype being changed by deepspeed for certain regression metrics ([#2379](https://github.com/Lightning-AI/torchmetrics/pull/2379))
 
 
+- Fix getitem for metric collection when prefix/postfix is set ([#2430](https://github.com/Lightning-AI/torchmetrics/pull/2430))
+
 ## [1.3.1] - 2024-02-12
 
 ### Fixed

src/torchmetrics/collections.py

@@ -547,6 +547,10 @@ def __getitem__(self, key: str, copy_state: bool = True) -> Metric:
 
  """
  self._compute_groups_create_state_ref(copy_state)
+ if self.prefix is not None:
+ key = key.removeprefix(self.prefix)
+ if self.postfix is not None:
+ key = key.removesuffix(self.postfix)
  return self._modules[key]
 
  @staticmethod

tests/unittests/bases/test_collections.py

@@ -204,6 +204,10 @@ def test_metric_collection_prefix_postfix_args(prefix, postfix):
  for name in names:
  assert f"new_prefix_{name}_new_postfix" in out, "postfix argument not working as intended with clone method"
 
+ keys = list(new_metric_collection.keys())
+ for k in keys:
+ assert new_metric_collection[k] # check that the keys are valid even with prefix and postfix
+
 
 def test_metric_collection_repr():
  """Test MetricCollection."""