Fix warning in test_trainer

Lightning-Universe · Apr 16, 2021 · f0faf2b · f0faf2b
1 parent 656d123
commit f0faf2b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/core/test_trainer.py b/tests/core/test_trainer.py
@@ -43,7 +43,7 @@ class DummyClassifier(nn.Module):
     def __init__(self):
         super().__init__()
         self.backbone = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10))
-        self.head = nn.LogSoftmax()
+        self.head = nn.LogSoftmax(dim=1)
 
     def forward(self, x):
         return self.head(self.backbone(x))