Fix issue #654 (#655)

Avoid two softmax calls in LogisticRegression Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lightning-Universe · Jun 16, 2021 · b087892 · b087892
1 parent 50969a8
commit b087892
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pl_bolts/models/regression/logistic_regression.py b/pl_bolts/models/regression/logistic_regression.py
@@ -53,7 +53,7 @@ def training_step(self, batch, batch_idx):
         # flatten any input
         x = x.view(x.size(0), -1)
 
-        y_hat = self(x)
+        y_hat = self.linear(x)
 
         # PyTorch cross_entropy function combines log_softmax and nll_loss in single function
         loss = F.cross_entropy(y_hat, y, reduction='sum')