Add test running predict inside the LightningModule hooks

[awaelchli]

🚀 Feature

Motivation

This should work. Add a test to make sure this works reliably in multi-gpu setting.

class MyModel(pl.LightningModule):
    def on_training_epoch_start():
        # calling trainer.predict inside trainer.fit routine.
        # is this okay?
        pred_outputs = self.trainer.predict(dataloaders=self.datamodule.train_loader)
        # do some work with pred_outputs
# start training
model = MyModel()
trainer = Trainer(model, accelerator='ddp')
trainer.fit()

Pitch

Add a test. Integrate findings in future loop design.

Alternatives

User can write their own predict loop inside the hooks, but is error prone and untested.

Additional context

Asked on slack by Dilip Thiagarajan