From c25645158364f053fe2a856ffb3753b474a7c5e2 Mon Sep 17 00:00:00 2001
From: David Landup <60978046+DavidLandup0@users.noreply.github.com>
Date: Sun, 4 Jun 2023 16:22:28 +0200
Subject: [PATCH] SegFormer MulticlassAccuracy metric fix (#139)

---
 .../segmentation/segformer/segformer_pt.py    | 37 ++++++++++---------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/deepvision/models/segmentation/segformer/segformer_pt.py b/deepvision/models/segmentation/segformer/segformer_pt.py
index 8d5bc1c..b8f4b8d 100644
--- a/deepvision/models/segmentation/segformer/segformer_pt.py
+++ b/deepvision/models/segmentation/segformer/segformer_pt.py
@@ -40,7 +40,8 @@ def __init__(
             backend="pytorch",
         )
         self.softmax_output = softmax_output
-        self.acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
+        if num_classes > 1:
+            self.acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
 
     def forward(self, x):
         y = self.backbone(x)
@@ -74,14 +75,15 @@ def training_step(self, train_batch, batch_idx):
             on_epoch=True,
             prog_bar=True,
         )
-        acc = self.acc(outputs, targets)
-        self.log(
-            "acc",
-            acc,
-            on_step=True,
-            on_epoch=True,
-            prog_bar=True,
-        )
+        if self.num_classes > 1:
+            acc = self.acc(outputs, targets)
+            self.log(
+                "acc",
+                acc,
+                on_step=True,
+                on_epoch=True,
+                prog_bar=True,
+            )
         return loss
 
     def validation_step(self, val_batch, batch_idx):
@@ -95,12 +97,13 @@ def validation_step(self, val_batch, batch_idx):
             on_epoch=True,
             prog_bar=True,
         )
-        val_acc = self.acc(outputs, targets)
-        self.log(
-            "val_acc",
-            val_acc,
-            on_step=True,
-            on_epoch=True,
-            prog_bar=True,
-        )
+        if self.num_classes > 1:
+            val_acc = self.acc(outputs, targets)
+            self.log(
+                "val_acc",
+                val_acc,
+                on_step=True,
+                on_epoch=True,
+                prog_bar=True,
+            )
         return loss