remove functional call from sigmoid and softmax (#255)

Lightning-Universe · Apr 29, 2021 · 2b39eae · 2b39eae
1 parent b275d8b
commit 2b39eae
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flash/core/classification.py b/flash/core/classification.py
@@ -60,8 +60,8 @@ def __init__(
 
     def to_metrics_format(self, x: torch.Tensor) -> torch.Tensor:
         if getattr(self.hparams, "multi_label", False):
-            return F.sigmoid(x)
-        return F.softmax(x, -1)
+            return torch.sigmoid(x)
+        return torch.softmax(x, -1)
 
 
 class ClassificationSerializer(Serializer):