Fix for dispatcher

Author: gianscarpe

tests/models/test_restore.py

@@ -25,7 +25,7 @@
 
 import tests.base.develop_pipelines as tpipes
 import tests.base.develop_utils as tutils
-from pytorch_lightning import Trainer, LightningModule, Callback, seed_everything
+from pytorch_lightning import Callback, LightningModule, Trainer, seed_everything
 from pytorch_lightning.callbacks import ModelCheckpoint
 from tests.base import EvalModelTemplate, GenericEvalModelTemplate, TrialMNIST
 
@@ -139,6 +139,8 @@ def test_callbacks_references_resume_from_checkpoint(tmpdir):
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 def test_running_test_pretrained_model_distrib_dp(tmpdir):
     """Verify `test()` on pretrained model."""
+    __import__("pdb").set_trace()
+
     tutils.set_random_master_port()
 
     model = EvalModelTemplate()
@@ -183,7 +185,7 @@ def test_running_test_pretrained_model_distrib_dp(tmpdir):
         dataloaders = [dataloaders]
 
     for dataloader in dataloaders:
-        tpipes.run_prediction(dataloader, pretrained_model)
+        tpipes.run_prediction(pretrained_model, dataloader)
 
 
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
@@ -234,7 +236,7 @@ def test_running_test_pretrained_model_distrib_ddp_spawn(tmpdir):
         dataloaders = [dataloaders]
 
     for dataloader in dataloaders:
-        tpipes.run_prediction(dataloader, pretrained_model)
+        tpipes.run_prediction(pretrained_model, dataloader)
 
 
 def test_running_test_pretrained_model_cpu(tmpdir):
@@ -376,7 +378,7 @@ def assert_good_acc():
         dp_model.eval()
 
         dataloader = trainer.train_dataloader
-        tpipes.run_prediction(dataloader, dp_model, dp=True)
+        tpipes.run_prediction(dp_model, dataloader, dp=True)
 
     # new model
     model = EvalModelTemplate(**hparams)