DataModule with lr_find not supported
🐛 Bug

To Reproduce
Steps to reproduce the behavior:

Create a Trainer (trainer), a LightningModule (model) and a DataModule (data_module)
Call trainer.lr_find(model, datamodule=data_module)
Error: TypeError: lr_find() got an unexpected keyword argument 'datamodule'


Code sample
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader
class DataModule(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
    def gen_set(self, start, end):
        vals = torch.tensor([float(i) for i in range(0, 800)])
        return torch.stack([torch.sin(vals), torch.cos(vals)], 1)
    def train_dataloader(self):
        return DataLoader(dataset=self.gen_set(0, 800))
    def val_dataloader(self):
        return DataLoader(dataset=self.gen_set(800, 900))
    def test_dataloader(self):
        return DataLoader(dataset=self.gen_set(900, 1000))

class LitModel(pl.LightningModule):
    def __init__(self, in_features):
        super().__init__()
        self.in_layer = nn.Linear(in_features, in_features//2)
        self.out_layer = nn.Linear(in_features//2, in_features)
        self.criterion = nn.MSELoss()
        self.lr = 0.001
    def forward(self, inp):
        inp = torch.tanh(self.in_layer(inp))
        return self.out_layer(inp)
    def training_step(self, batch, batch_idx):
        x = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, x.float())
        return pl.TrainResult(minimize=loss)
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=(self.lr or self.learning_rate))
    
data_module = DataModule()
trainer = pl.Trainer()
model = LitModel(in_features=2)
lr_finder = trainer.lr_find(model, datamodule=data_module)   #error

Expected behavior

The lr_find() function should work as if it were passed a train_loader and test_loader when invoking it, just like fit()
Environment

Packages:

numpy:             1.18.5
pyTorch_debug:     False
pyTorch_version:   1.5.1
pytorch-lightning: 0.9.0
tensorboard:       2.2.0
tqdm:              4.45.0


System:

OS:                Linux
architecture:

64bit


processor:         x86_64
python:            3.7.6