change for log softmax

Lightning-Universe · Apr 21, 2021 · 2fed6a2 · 2fed6a2
1 parent 644f062
commit 2fed6a2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
@@ -168,7 +168,7 @@ class Foo(ImageClassifier):
 
 def test_optimization(tmpdir):
 
-    model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.Softmax())
+    model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax())
     optim = torch.optim.Adam(model.parameters())
     task = ClassificationTask(model, optimizer=optim, scheduler=None)