Fix PyTorch Lightning dev version compatibility (mlflow#4959)

* Fix Signed-off-by: dbczumar <corey.zumar@databricks.com> * Add a newline Signed-off-by: dbczumar <corey.zumar@databricks.com> * Remove space :) Signed-off-by: westford14 <westford14@gmail.com>
westford14 · Nov 2, 2021 · 76b8a39 · 76b8a39
1 parent e75b47c
commit 76b8a39
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/pytorch/test_pytorch_autolog.py b/tests/pytorch/test_pytorch_autolog.py
@@ -110,7 +110,7 @@ def test_pytorch_autolog_persists_manually_created_run():
         dm.setup(stage="fit")
         trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
         trainer.fit(model, dm)
-        trainer.test()
+        trainer.test(datamodule=dm)
         assert mlflow.active_run() is not None
         assert mlflow.active_run().info.run_id == manual_run.info.run_id