Cannot pickle custom metric with DDP mode.
🐛 Bug

To Reproduce
Steps to reproduce the behavior:

just run this script.


Code sample

import torch
import torch.nn.functional as F
import pytorch_lightning as pl
from pytorch_lightning.metrics import Metric, TensorMetric

class MetricPerplexity(Metric):
    """
    Computes the perplexity of the model.
    """

    def __init__(self, pad_idx: int, *args, **kwargs):
        super().__init__(name='ppl')
        self.pad_idx = pad_idx

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        
        loss = F.cross_entropy(pred, target, reduction='none')
        non_padding = target.ne(self.pad_idx)
        loss = loss.masked_select(non_padding).sum()

        num_words = non_padding.sum()
        ppl = torch.exp(
            torch.min(loss / num_words, torch.tensor([100]).type_as(loss))
        )
        return ppl

class TensorPerplexity(TensorMetric):
    """
    Computes the perplexity of the model.
    """

    def __init__(self, pad_idx: int, *args, **kwargs):
        super().__init__(name='ppl', *args, **kwargs)
        self.pad_idx = pad_idx

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        
        loss = F.cross_entropy(pred, target, reduction='none')
        non_padding = target.ne(self.pad_idx)
        loss = loss.masked_select(non_padding).sum()

        num_words = non_padding.sum()
        ppl = torch.exp(
            torch.min(loss / num_words, torch.tensor([100]).type_as(loss))
        )
        return ppl

class ModelWithMetric(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(50, 1)
        self.ppl = MetricPerplexity(100)
    
    def training_step(self, batch, batch_nb):
        return {'loss': torch.mean(self(batch))}
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), 0.01)
    
    def forward(self, batch):
        return self.lin(batch[0])

class ModelWithTensorMetric(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(50, 1)
        self.ppl = TensorPerplexity(100)
    
    def training_step(self, batch, batch_nb):
        return {'loss': torch.mean(self(batch))}
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), 0.01)
    
    def forward(self, batch):
        return self.lin(batch[0])

if __name__ == "__main__":
    m1 = ModelWithMetric()
    m2 = ModelWithTensorMetric()

    t1 = pl.Trainer(distributed_backend='ddp_cpu', fast_dev_run=True)
    t2 = pl.Trainer(distributed_backend='ddp_cpu', fast_dev_run=True)

    loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(torch.rand(10, 50)), batch_size=5)
    t1.fit(m1, train_dataloader=loader)
    print('works well')
    t2.fit(m2, train_dataloader=loader)

Error log
AttributeError: Can't pickle local object '_apply_to_outputs.<locals>.decorator_fn.<locals>.new_func'

Expected behavior
I do not know how to explain this behaviour, I think both of this classes should work well with DDP.

Environment


OS: MacOS


CUDA:
- GPU:
- available:         False
- version:           None


Packages:
- numpy:             1.18.5
- pyTorch_debug:     False
- pyTorch_version:   1.6.0
- pytorch-lightning: 0.9.0rc9
- tensorboard:       2.2.2
- tqdm:              4.48.0


System:
- OS:                Darwin
- architecture:
- 64bit
- processor:         i386
- python:            3.7.7
- version:           Darwin Kernel Version 19.6.0: Sun Jul  5 00:43:10 PDT 2020


Additional context