unpin torchmetrics

bmfm_targets/training/losses/task.py

@@ -197,13 +197,18 @@ def extract_metric_inputs(
 
         return self.metric_key, model_outputs, gt_labels
 
-    def get_metrics(self) -> MetricCollection:
+    def get_metrics(self, exclude: set[str] | None = None) -> MetricCollection:
         """
         Get metric collection for this task.
 
         Uses objective's default metrics unless overridden via metrics parameter.
         Filters metrics based on output size (classification vs regression).
 
+        Parameters
+        ----------
+        exclude:
+            Optional set of metric names to omit from the collection.
+
         Returns
         -------
             MetricCollection: Collection of metrics for this task
@@ -227,6 +232,7 @@ def get_metrics(self) -> MetricCollection:
             {
                 mt["name"]: metrics.get_metric_object(mt, num_classes)
                 for mt in metric_configs
+                if not (exclude and mt["name"] in exclude)
             }
         )
 

bmfm_targets/training/modules/base.py

@@ -249,12 +249,11 @@ def initialize_metrics(self):
 
         metrics_dict = {}
         for loss_task in self.loss_tasks:
-            metrics_dict[loss_task.metric_key] = loss_task.get_metrics()
-            # Remove perplexity when multiple tasks share a metric_key
-            # (perplexity needs logits, but duplicated keys use predictions)
-            if metric_key_counts[loss_task.metric_key] > 1:
-                if "perplexity" in metrics_dict[loss_task.metric_key]:
-                    metrics_dict[loss_task.metric_key].pop("perplexity")
+            # Perplexity needs logits, but duplicated metric_keys use predictions
+            exclude = (
+                {"perplexity"} if metric_key_counts[loss_task.metric_key] > 1 else None
+            )
+            metrics_dict[loss_task.metric_key] = loss_task.get_metrics(exclude=exclude)
 
         self.train_metrics = MultitaskWrapper(metrics_dict).clone()
         self.val_metrics = MultitaskWrapper(metrics_dict).clone()

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
   "hydra-core",
   "clearml>1.13,<2",
   "rich",
-  "torchmetrics==1.1.0",
+  "torchmetrics",
   "tensorboardX",
   "pandas>=2,<3",
   "einops",