Fix test

Lightning-Universe · Apr 24, 2021 · 63839f5 · 63839f5
1 parent 67f49bd
commit 63839f5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/core/test_trainer.py b/tests/core/test_trainer.py
@@ -53,7 +53,7 @@ def test_task_fit(tmpdir: str):
     model = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10), nn.LogSoftmax())
     train_dl = torch.utils.data.DataLoader(DummyDataset())
     val_dl = torch.utils.data.DataLoader(DummyDataset())
-    task = ClassificationTask(model, F.nll_loss)
+    task = ClassificationTask(model, loss_fn=F.nll_loss)
     trainer = Trainer(fast_dev_run=True, default_root_dir=tmpdir)
     result = trainer.fit(task, train_dl, val_dl)
     assert result
@@ -63,7 +63,7 @@ def test_task_finetune(tmpdir: str):
     model = DummyClassifier()
     train_dl = torch.utils.data.DataLoader(DummyDataset())
     val_dl = torch.utils.data.DataLoader(DummyDataset())
-    task = ClassificationTask(model, F.nll_loss)
+    task = ClassificationTask(model, loss_fn=F.nll_loss)
     trainer = Trainer(fast_dev_run=True, default_root_dir=tmpdir)
     result = trainer.finetune(task, train_dl, val_dl, strategy=NoFreeze())
     assert result