Update tests for SemSegment

Lightning-Universe · Nov 3, 2020 · 037434a · 037434a
1 parent 9809624
commit 037434a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/models/test_vision.py b/tests/models/test_vision.py
@@ -67,10 +67,10 @@ def train_dataloader(self):
 
  dm = DummyDataModule()
 
- model = SemSegment(datamodule=dm, num_classes=19)
+ model = SemSegment(num_classes=19)
 
  trainer = pl.Trainer(fast_dev_run=True, max_epochs=1)
- trainer.fit(model)
+ trainer.fit(model, dm)
  loss = trainer.progress_bar_dict['loss']
 
  assert float(loss) > 0