def METHOD_NAME(tmpdir, datamodule):
    class TestModel(LightningModule):
        def __init__(self):
            super().__init__()
            self.layer = torch.nn.Linear(32, 2)
        def forward(self, x):
            return self.layer(x)
    class TestLightningDataModule(LightningDataModule):
        def __init__(self, dataloaders):
            super().__init__()
            self._dataloaders = dataloaders
        def test_dataloader(self):
            return self._dataloaders
        def predict_dataloader(self):
            return self._dataloaders
    data = [torch.utils.data.DataLoader(RandomDataset(32, 2)), torch.utils.data.DataLoader(RandomDataset(32, 2))]
    if datamodule:
        data = TestLightningDataModule(data)
    model = TestModel()
    trainer = Trainer(default_root_dir=tmpdir)
    results = trainer.predict(model, data)
    assert len(results) == 2
    assert results[0][0].shape == torch.Size([1, 2])