Feat: make __getattr__ and __setattr__ of ClasswiseWrapper more general

CHANGELOG.md

@@ -20,7 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
--
+- Made `__getattr__` and `__setattr__` of `ClasswiseWrapper` more general ([#2424](https://github.com/Lightning-AI/torchmetrics/pull/2424))
 
 
 ### Deprecated

src/torchmetrics/wrappers/classwise.py

@@ -216,19 +216,19 @@ def plot(
 
     def __getattr__(self, name: str) -> Union[Tensor, "Module"]:
         """Get attribute from classwise wrapper."""
-        # return state from self.metric
-        if name in ["tp", "fp", "fn", "tn"]:
-            return getattr(self.metric, name)
+        if name == "metric" or (name in self.__dict__ and name not in self.metric.__dict__):
+            # we need this to prevent from infinite getattribute loop.
+            return super().__getattr__(name)
 
-        return super().__getattr__(name)
+        return getattr(self.metric, name)
 
     def __setattr__(self, name: str, value: Any) -> None:
         """Set attribute to classwise wrapper."""
-        super().__setattr__(name, value)
-        if name == "metric":
-            self._defaults = self.metric._defaults
-            self._persistent = self.metric._persistent
-            self._reductions = self.metric._reductions
-        if hasattr(self, "metric") and name in ["tp", "fp", "fn", "tn", "_update_count", "_computed"]:
-            # update ``_update_count`` and ``_computed`` of internal metric to prevent warning.
+        if hasattr(self, "metric") and name in self.metric._defaults:
             setattr(self.metric, name, value)
+        else:
+            super().__setattr__(name, value)
+            if name == "metric":
+                self._defaults = self.metric._defaults
+                self._persistent = self.metric._persistent
+                self._reductions = self.metric._reductions