fixed out of bounds for auroc computation on betti numbers

aidos-lab · Aug 16, 2024 · 65cd162 · 65cd162
1 parent c8f76b5
commit 65cd162
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/code/metrics/custom_metrics/auroc.py b/code/metrics/custom_metrics/auroc.py
@@ -17,6 +17,10 @@ def __init__(self, num_classes: int, **kwargs):
         )
 
     def update(self, preds: Tensor, target: Tensor) -> None:
+        preds = torch.min(
+            torch.max(preds, torch.tensor(0.0)),
+            torch.tensor(self.num_classes - 1),
+        ).long()
         y_hat = torch.nn.functional.one_hot(preds, self.num_classes).float()
         self.auroc.update(preds=y_hat, target=target)