Strange issue with DataParallel
🐛 Bug

I'm having an issue with tensors being on different devices when using distributed_backend=dp. But it only occurs under what seems like some pretty specific circumstances. I can only reproduce the bug when I have done BOTH of the following:

replaced the forward method of the internal network (see the code sample below)
replaced the BatchSampler in the training dataloader with my own version

Everything is fine if I do only one of those two things. And I don't have any problem doing both of them if I train on a single GPU. What's doubly strange is that I only replace the BatchSampler for the train loader, but I encounter the error during the initial validation sanity check. I really have no idea what's going on here.
To Reproduce
Here's code to reproduce. Sorry it's a little long.
import os
from collections import defaultdict
import random
import torch
from torch.utils.data import Dataset
from pytorch_lightning import Trainer, LightningModule
import torchvision


class BalancedBatchSampler(torch.utils.data.sampler.BatchSampler):
    '''Samples a batch by first randomly choosing n_classes classes and then
    n_samples images from each class'''
    def __init__(self, dataset, batch_size, n_samples):
        self.targets = dataset.label
        self.target_idx = defaultdict(list)
        for i, t in enumerate(self.targets):
            self.target_idx[t].append(i)
        for ii in self.target_idx.values():
            random.shuffle(ii)
        self.used_count = {t: 0 for t in self.target_idx}
        self.n_classes = batch_size // n_samples
        self.n_samples = n_samples
        self.dataset = dataset
        self.batch_size = self.n_samples * self.n_classes

    def __iter__(self):
        self.count = 0
        nsamp = self.n_samples
        while self.count + self.batch_size < len(self.dataset):
            classes = random.sample(self.target_idx.keys(), self.n_classes)
            indices = []
            for c in classes:
                used = self.used_count[c]
                indices.extend(self.target_idx[c][used:used+nsamp])
                self.used_count[c] += nsamp
                if self.used_count[c] + nsamp > len(self.target_idx[c]):
                    random.shuffle(self.target_idx[c])
                    self.used_count[c] = 0
            yield indices
            self.count += self.n_classes * self.n_samples

    def __len__(self):
        return len(self.dataset) // self.batch_size


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, *size)
        self.label = torch.randint(0,10,(length,))

    def __getitem__(self, index):
        return self.data[index], self.label[index]

    def __len__(self):
        return self.len

def ResnetBase(resnet_model, forward_func=None):
    '''Return a torchvision.models.ResNet model without the average pooling
    and fully connected layers at the end. An optional forward_func can be
    passed to redefine the forward pass of the model (for instance, to return
    intermediate layer outputs).'''
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x

    if forward_func is None: forward_func = forward

    cout = resnet_model.fc.in_features
    resnet_model.cout = cout
    resnet_model.forward = forward_func.__get__(resnet_model)

    delattr(resnet_model, 'avgpool')
    delattr(resnet_model, 'fc')

    return resnet_model


class BoringModel(LightningModule):

    def __init__(self, layer):
        """
        Testing PL Module

        Use as follows:
        - subclass
        - modify the behavior for what you want

        class TestModel(BaseTestModel):
            def training_step(...):
                # do your own thing

        or:

        model = BaseTestModel()
        model.training_epoch_end = None

        """
        super().__init__()
        net = torchvision.models.resnet18(pretrained=False)
        #self.net = net  # This works fine
        self.net = ResnetBase(net)  # This does not
        self.layer = torch.nn.Conv2d(3, 32, 3)
        self.use_layer = layer

    def forward(self, x):
        if self.use_layer:
            x = self.layer(x)
        else:
            x = self.net(x)
        return x

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))

    def training_step(self, batch, batch_idx):
        x, y = batch
        output = self(x)
        loss = self.loss(batch, output)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        output = self(x)
        loss = self.loss(batch, output)
        return {"x": loss}

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]


def run_test(args):
    # fake data
    train_data = RandomDataset((3,64,64), 128)
    val_data = RandomDataset((3,64,64), 128)
    if args.sampler:
        sampler = BalancedBatchSampler(train_data, 32, 4)
    else:
        sampler = None
    train_data = torch.utils.data.DataLoader(train_data, batch_sampler=sampler)
    val_data = torch.utils.data.DataLoader(val_data)

    # model
    model = BoringModel(args.layer)
    trainer = Trainer.from_argparse_args(
        args, 
        default_root_dir=os.getcwd(),
        limit_train_batches=5,
        limit_val_batches=5,
        max_epochs=1,
        weights_summary=None,
        distributed_backend='dp',
        gpus=2,
    )
    trainer.fit(model, train_data, val_data)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--sampler', action='store_true')
    parser.add_argument('--layer', action='store_true')
    parser = Trainer.add_argparse_args(parser)
    args = parser.parse_args()
    run_test(args)
If you save the above in bug.py and then run
python bug.py --sampler

You should get the following error
Traceback (most recent call last):
  File "dpbug.py", line 195, in <module>
    run_test(args)
  File "dpbug.py", line 185, in run_test
    trainer.fit(model, train_data, val_data)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 440, in fit
    results = self.accelerator_backend.train()
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/accelerators/dp_accelerator.py", line 97, in train
    results = self.train_or_test()
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 68, in train_or_test
    results = self.trainer.train()
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 485, in train
    self.train_loop.run_training_epoch()
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 544, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 713, in run_training_batch
    self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 452, in optimizer_step
    self.trainer.accelerator_backend.optimizer_step(
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 115, in optimizer_step
    model_ref.optimizer_step(
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1216, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
    return wrapped(*args, **kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
    return func(*args, **kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/optim/sgd.py", line 86, in step
    loss = closure()
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 703, in train_step_and_backward_closure
    result = self.training_step_and_backward(
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 798, in training_step_and_backward
    result = self.training_step(split_batch, batch_idx, opt_idx, hiddens)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 305, in training_step
    training_step_output = self.trainer.accelerator_backend.training_step(args)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/accelerators/dp_accelerator.py", line 111, in training_step
    output = self.trainer.model(*args)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/overrides/data_parallel.py", line 87, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/overrides/data_parallel.py", line 151, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/overrides/data_parallel.py", line 310, in parallel_apply
    raise output
  File "/home/catalys1/pylt/lib/python3.8/site-packages/pytorch_lightning/overrides/data_parallel.py", line 263, in _worker
    output = module.training_step(*input, **kwargs)
  File "dpbug.py", line 147, in training_step
    output = self(x)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "dpbug.py", line 138, in forward
    x = self.net(x)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "dpbug.py", line 85, in forward
    x = self.conv1(x)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 419, in forward
    return self._conv_forward(input, self.weight)
  File "/home/catalys1/pylt/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 415, in _conv_forward
    return F.conv2d(input, weight, self.bias, self.stride,
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 2 does not equal 1 (while checking arguments for cudnn_convolution)

Running it as either
python bug.py --sampler --layer

or
python bug.py

gives no error
I'm on lightning version 1.0.4.
Expected behavior
Normal operation without strange device mismatch error.
Environment

PyTorch Version (e.g., 1.0): 1.0.4
OS (e.g., Linux): Linux
How you installed PyTorch (conda, pip, source): pip
Python version: 3.8
CUDA/cuDNN version: 10.2
GPU models and configuration: 1080 Ti