ADD: threshold

interTwin-eu · Nov 7, 2023 · 55aac61 · 55aac61
1 parent be865f3
commit 55aac61
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/itwinai/torch/inference.py b/src/itwinai/torch/inference.py
@@ -187,6 +187,8 @@ class MultilabelTorchPredictor(TorchPredictor):
     output of the neural network.
     """
 
+    threshold: float
+
     def __init__(
         self,
         model: Union[nn.Module, ModelLoader],
@@ -200,6 +202,9 @@ def __init__(
         )
         self.threshold = threshold
 
+    def transform_predictions(self, batch: Batch) -> Batch:
+        return (batch > self.threshold).float()
+
 
 class RegressionTorchPredictor(TorchPredictor):
     """