From 2fed6a29f4deea33329fe59f39b086a99445c529 Mon Sep 17 00:00:00 2001 From: tchaton Date: Wed, 21 Apr 2021 11:11:52 +0100 Subject: [PATCH] change for log softmax --- tests/core/test_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core/test_model.py b/tests/core/test_model.py index 5b18a711db..2426cff78f 100644 --- 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)