I am trying to use pytorch_lightning with multiple GPU, but get the following error:
RuntimeError: All input tensors must be on the same device. Received cuda:0 and cuda:3

How to fix this? Below is a MWE:
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import pytorch_lightning as pl

class DataModule(pl.LightningDataModule):

    def __init__(self):
        super().__init__()
    
    def setup(self, stage):
        # called on each gpu
        n = 128
        x = torch.randn(n, 1, 64,64)
        data = list(zip(x,x))
        self.test  = DataLoader(data, batch_size=32)
        self.train = DataLoader(data, batch_size=32)
        self.val   = DataLoader(data, batch_size=32)
        
    def train_dataloader(self):
        return self.train

    def val_dataloader(self):
        return self.val

    def test_dataloader(self):
        return self.test
    
class Net(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.net = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(1,1))
    
    def forward(self, x):
        return self.net(x)
    
    def validation_step(self, batch, batch_idx):
        loss = self.compute_batch_loss(batch, batch_idx)
        self.log('val_loss', loss)
        return loss
    
    def compute_batch_loss(self, batch, batch_idx):
        x, y = batch
        y_hat = self.net(x)
        loss = F.mse_loss(y_hat, y)
        return loss

    def training_step(self, batch, batch_idx):
        loss = self.compute_batch_loss(batch, batch_idx)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

dm = DataModule()
model = Net()
trainer = pl.Trainer(gpus=4, 
                     distributed_backend="dp",
                     max_epochs=1,
                    )
trainer.fit(model, dm)