From 2b39eae8e9d6efa40ed9a8ffb57e08b269b418f4 Mon Sep 17 00:00:00 2001
From: Edgar Riba <edgar.riba@gmail.com>
Date: Thu, 29 Apr 2021 18:02:40 +0200
Subject: [PATCH] remove functional call from sigmoid and softmax (#255)

---
 flash/core/classification.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flash/core/classification.py b/flash/core/classification.py
index 346905b823..b63a9c8b58 100644
--- 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):