Passing dataloader to trainer.fit() doesn't work with tpu (and maybe ddp)
🐛 Bug
Receive a
AttributeError: Can't pickle local object 'Trainer.__set_fit_dataloaders.<locals>.patch_train_dataloader'

error when passing the dataloader directly to trainer.fit(model, train_loader)
To Reproduce
Steps to reproduce the behavior:
Try to call trainer.fit(model, train_loader) in TPU mode.
(I suspect that anything that calls mp.spawn will cause this problem, so ddp probably will face this issue too.)

Code sample
import os

import pytorch_lightning as pl
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

import torch_xla.core.xla_model as xm


class CoolSystem(pl.LightningModule):
    def __init__(self, use_tpu=False):
        super(CoolSystem, self).__init__()
        # not the best model...
        self.use_tpu = use_tpu
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_idx):
        # REQUIRED
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)
        return {'val_loss': F.cross_entropy(y_hat, y)}

    def validation_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        tensorboard_logs = {'val_loss': avg_loss}
        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}

    def test_step(self, batch, batch_idx):
        # OPTIONAL
        x, y = batch
        y_hat = self.forward(x)
        return {'test_loss': F.cross_entropy(y_hat, y)}

    def test_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
        tensorboard_logs = {'test_loss': avg_loss}
        return {'avg_test_loss': avg_loss, 'log': tensorboard_logs}

    def configure_optimizers(self):
        # REQUIRED
        # can return multiple optimizers and learning_rate schedulers
        # (LBFGS it is automatically supported, no need for closure function)
        return torch.optim.Adam(self.parameters(), lr=0.0004)


if __name__ == '__main__':
    from pytorch_lightning import Trainer

    model = CoolSystem(use_tpu=True)

    dataset = MNIST(os.getcwd(),
                        train=True,
                        download=True,
                        transform=transforms.ToTensor())

    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=True)

    loader = DataLoader(dataset, sampler=sampler, batch_size=32)

    # most basic trainer, uses good defaults
    trainer = Trainer(num_tpu_cores=8)
    trainer.fit(model, loader)
Expected behavior

Ideally, specifying the dataloaders as part of the LightningModule should work just the same as passing the dataloaders into trainer.fit()
Environment
Please copy and paste the output from our
environment collection script
(or fill out the checklist below manually).
You can get the script and run it with:
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py


Docker image: gcr.io/tpu-pytorch/xla:nightly
build steps:

pip install git+git://github.com/williamFalcon/pytorch-lightning.git@master --upgrade



(I don't have access to the machine right now, so please forgive me on the specific version info temporarily)
Proposed solution
The issue is here, trying to assign a local function to the model
def __set_fit_dataloaders(self, model, train_dataloader, val_dataloaders, test_dataloaders):
        # when dataloader is passed via fit, patch the train_dataloader
        # functions to overwrite with these implementations
        if train_dataloader is not None:
            if not self.is_overriden('training_step', model):
                m = 'You called .fit() with a train_dataloader but did not define training_step()'
                raise MisconfigurationException(m)

            def patch_train_dataloader():
                return train_dataloader

            model.train_dataloader = patch_train_dataloader
Instead of using a closure or a local function, you could use a callable defined at the top-level.  This will be pickleable.
class DataLoaderPatcher
    def __init__(self, loader):
        self.loader = loader
    def __call__():
        return self.loader

def __set_fit_dataloaders(self, model, train_dataloader, val_dataloaders, test_dataloaders):
        # when dataloader is passed via fit, patch the train_dataloader
        # functions to overwrite with these implementations
        if train_dataloader is not None:
            if not self.is_overriden('training_step', model):
                m = 'You called .fit() with a train_dataloader but did not define training_step()'
                raise MisconfigurationException(m)

            model.train_dataloader = DataLoaderPatcher(train_dataloader)