Add task arg to torchmetrics metrics in relevant examples

GRAAL-Research · Dec 2, 2022 · bab2195 · bab2195
1 parent 1130501
commit bab2195
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/docs/source/examples/semantic_segmentation.rst b/docs/source/examples/semantic_segmentation.rst
@@ -252,7 +252,7 @@ For training, we use the Jaccard index metric in addition to the accuracy and F1
         optimizer,
         criterion,
         batch_metrics=['accuracy'],
-        epoch_metrics=['f1', torchmetrics.JaccardIndex(num_classes=22)],
+        epoch_metrics=['f1', torchmetrics.JaccardIndex(num_classes=22, task="multiclass")],
         device=device,
     )
 

diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst
@@ -173,22 +173,22 @@ Examples:
 
     metric_collection = MetricCollection(
         [
-            F1Score(num_classes=10, average='macro'),
-            Precision(num_classes=10, average='macro'),
-            Recall(num_classes=10, average='macro'),
+            F1Score(num_classes=10, average="macro", task="multiclass"),
+            Precision(num_classes=10, average="macro", task="multiclass"),
+            Recall(num_classes=10, average="macro", task="multiclass"),
         ]
     )
 
     metrics = [
-        ('custom_name', my_custom_metric),
-        (('metric_1', 'metric_2'), my_custom_metric2),
-        (('a', 'b'), my_custom_metric3),
-        (('metric_3', 'metric_4'), CustomMetric()),
-        (('c', 'd'), CustomMetric2()),
+        ("custom_name", my_custom_metric),
+        (("metric_1", "metric_2"), my_custom_metric2),
+        (("a", "b"), my_custom_metric3),
+        (("metric_3", "metric_4"), CustomMetric()),
+        (("c", "d"), CustomMetric2()),
 
         # No need to pass the names since the class sets the attribute __name__.
         CustomMetric3(),
 
         # The names are the keys returned by MetricCollection.
-        (('F1Score', 'Precision', 'Recall'), metric_collection),
+        (("F1Score", "Precision", "Recall"), metric_collection),
     ]
diff --git a/examples/semantic_segmentation.ipynb b/examples/semantic_segmentation.ipynb
@@ -990,7 +990,7 @@
     "    optimizer,\n",
     "    criterion,\n",
     "    batch_metrics=['accuracy'],\n",
-    "    epoch_metrics=['f1', torchmetrics.JaccardIndex(num_classes=22)],\n",
+    "    epoch_metrics=['f1', torchmetrics.JaccardIndex(num_classes=22, task=\"multiclass\")],\n",
     "    device=device,\n",
     ")\n",
     "\n",