Fix incomplete test data module setup in Lightning inference

CHANGELOG.md

@@ -54,6 +54,7 @@ multiple large checkpoints can time out.
 mounting and running matplotblib on some machines. Re-instantiated a disabled test.
 - ([#509](https://github.com/microsoft/InnerEye-DeepLearning/pull/509)) Fix issue where model checkpoints were not loaded
 in inference-only runs when using lightning containers.
+- ([#553](https://github.com/microsoft/InnerEye-DeepLearning/pull/553)) Fix incomplete test data module setup in Lightning inference.
 
 ### Removed
 

InnerEye/ML/run_ml.py

@@ -510,7 +510,7 @@ def run_inference_for_lightning_models(self, checkpoint_paths: List[Path]) -> No
             # files to the right folder. Best guess is to change the current working directory to where files should go.
             with change_working_directory(self.container.outputs_folder):
                 trainer.test(self.container.model,
-                             test_dataloaders=self.container.get_data_module().test_dataloader())
+                             datamodule=self.container.get_data_module())
         else:
             logging.warning("None of the suitable test methods is overridden. Skipping inference completely.")