Batch is not split in 'dp' mode when dataloader output is not a tensor
My dataloader is returning a list of lists (for multi-label classification) for labels and a tensor of images for each batch. When I'm using DataParallel mode, labels are not getting split into "sub-batches" and I'm getting all the labels on each GPU. Is there a way to implement this splitting also for non-tensors?
class CustomDataset(Dataset):
...
    def collate(self, batch):
        images, labels = list(zip(*batch))
        return torch.stack(images), [ label_set for label_set in labels ]

class LitModel(pl.LightningModule):
...
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        print(f'Val loss {y.shape}, {y_hat.shape}' )

This example will print y to have length of full batch and y_hat to have length the same as the length of x tensor (which is smaller here). x tensor got split correctly, while y didn't.
Is it the issue of lightning or maybe DataParallel module?