Undefined variable on DDP2 trainer
🐛 Bug

I am trying to run an experiment in multiple gpus, but a single machine. If I understood right, using ddp2 would be faster than dp and would allow me to use the aggregated results in training_step_end and validation_step_end, which ddp doesn't provide. For some reason I am getting this error:
Traceback (most recent call last):
  File "main.py", line 50, in <module>
    trainer.fit(model)
  File "/home/darley.barreto/.local/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 882, in fit
    self.ddp_train(task, model)
UnboundLocalError: local variable 'task' referenced before assignment
Which is this line.
To Reproduce
I can't paste my code due to size and complexity, so I tweaked this code a bit and tried to run with ddp2, which gives the same error, but it works for ddp.
import os
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import torchvision.transforms as transforms

import pytorch_lightning as ptl
from pytorch_lightning import Trainer

class CoolModel(ptl.LightningModule):

    def __init__(self):
        super(CoolModel, self).__init__()
        # not the best model...
        self.l1 = torch.nn.Linear(28 * 28, 10)

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

    def my_loss(self, y_hat, y):
        return F.cross_entropy(y_hat, y)

    def training_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self.forward(x)
        return {'loss': self.my_loss(y_hat, y)}

    def validation_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self.forward(x)
        return {'val_loss': self.my_loss(y_hat, y)}

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

    def configure_optimizers(self):
        return [torch.optim.Adam(self.parameters(), lr=0.02)]

    def train_dataloader(self):
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

    def val_dataloader(self):
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

model = CoolModel()

trainer = Trainer(max_epochs=1, gpus=2, distributed_backend="ddp2", train_percent_check=0.1)
trainer.fit(model)
Environment
* CUDA:
        - GPU:
                - Quadro RTX 5000
                - Quadro RTX 5000
        - available:         True
        - version:           10.1
* Packages:
        - numpy:             1.18.1
        - pyTorch_debug:     False
        - pyTorch_version:   1.5.1
        - pytorch-lightning: 0.8.2-dev
        - tensorboard:       2.2.2
        - tqdm:              4.46.1
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - 
        - processor:         x86_64
        - python:            3.7.7
        - version:           #31-Ubuntu SMP Tue Jul 17 15:39:52 UTC 2018