Expected all tensors to be on the same device, but found at least two devices...

[hensel-f]

Hi, I'm very new to pytorch and pytorch lightning, and I'm uncertain how different devices (e.g. CPU/GPU) need to be handled. I'm want to train a network in such a way that after a fixed number of epochs there is an additional term (defined by a custom function) added to the loss function. However, as soon as the first epoch where the custom loss term occurs is reached I get the following error:

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat1 in method wrapper_addmm)

What am I doing wrong and how can I fix this issue?
Any help and insight is very much appreciated, thanks in advance!
If you have any clarifying questions please let me know.

Here is the model definition (except for the custom loss functions):

class Net_lightning(pl.LightningModule):
    def __init__(self, num_classes=10):
        super(Net_lightning, self).__init__()
        self.num_classes = num_classes
        self.flatten = nn.Flatten(-3, -1) 
        self.dense1 = nn.Linear(28 * 28, 64)
        self.act1 = nn.ReLU()
        self.dense2 = nn.Linear(64, 16) 
        self.act2 = nn.ReLU()
        self.dense3 = nn.Linear(16, self.num_classes) 
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.flatten(x)
        x = self.dense1(x)
        x = self.act1(x)
        x = self.dense2(x)
        x = self.act2(x)
        x = self.dense3(x)
        x = self.softmax(x)

        return x
        
        
def data_per_class_dl(dataloader):
    '''
    Given a pytorch Dataloader - initialized with an MPerClassSampler, to get the same number of samples
    per class in each batch - selects a random batch and groups it w.r.t. the target classes.

    :param dataloader: pytorch Dataloader
    :return: a tuple containing a list of tensors of the data part only of the data samples grouped
             by their label, and a sorted list of all labels occurring in the dataset.
    '''
    # choose a random batch containing a fixed number of representatives of each label:
    opt_batch = random.choice(list(dataloader))

    opt_batch_sorted = sorted(list(zip(torch.unbind(opt_batch[0], dim=0), torch.unbind(opt_batch[1], dim=0))),
                              key=lambda x: x[1])
    labels_unique = torch.unique(opt_batch[1], sorted=True)

    data_all_classes = []
    for t in labels_unique:
        data_class = torch.stack([d for (d, label) in opt_batch_sorted if label == t])
        data_all_classes.append(data_class)

    return (data_all_classes, labels_unique)

class c_optimizer_lightning(pl.LightningModule):
    '''
    Pytorch lightning module that trains a provided model (pl.Module) with custom-optimization after 'epoch_thresh' is reached.
    '''
    def __init__(self, model, learning_rate, loss_criterion, epoch_thresh, custom_loss_function, layers_id,
                 comparison_dataloader, learning_rate_ratio=1.0,matrix_norm_flag=False):
        '''
        :param model: The model (pl.Module) that will be trained.
        :param loss_criterion: The loss functioned used for training, e.g. torch.nn.CrossEntropyLoss().
        :param epoch_thresh: The threshold after which custom optimization kicks in (int).
        :param custom_loss_function: Function computing the custom loss.
        :param comparison_dataloader: A torch.utils.data.Dataloader that contains batches of class-balanced samples.
        :param layers_id: List of layers; List(int).
        :param comparison_balanced_selection: Selects a random batch of the comparison_dataloader and brings it into
                                              the required form to compute the custom loss.
        '''
        super(c_optimizer_lightning, self).__init__()
        self.model = model
        self.learning_rate = learning_rate
        self.loss_criterion = loss_criterion
        self.epoch_thresh = epoch_thresh
        self.custom_loss_function = custom_loss_function
        self.layers_id = layers_id
        self.comparison_dataloader = comparison_dataloader
        self.matrix_norm_flag = matrix_norm_flag
        self.learning_rate_ratio = learning_rate_ratio

    def training_step(self, batch, batch_idx):
        data, target = batch
        output = self.model(data)
        loss = self.loss_criterion(output, target)
        loss_c = 0
        if self.current_epoch >= self.epoch_thresh:
            all_classes_loss = False
            train_data_labels_per_cl = data_per_class_dl(self.comparison_dataloader)
            custom_step = self.custom_loss_function(self.model, list(zip(data, target)), train_data_labels_per_cl,
                                               layers_id=self.layers_id, all_classes_loss=all_classes_loss,
                                               matrix_norm_only=self.matrix_norm_flag)
            loss_c = torch.mean(custom_step)
            loss += self.learning_rate_ratio * loss_c
            self.log('custom loss', loss_c)

        preds = torch.argmax(output, dim=1)
        acc = accuracy(preds, target)
        self.log('loss', loss, prog_bar=True)
        self.log('acc', acc, prog_bar=True)

        return {'loss': loss, 'custom loss': loss_c, 'acc': acc}

    def evaluate(self, batch, stage=None):
        data, target = batch
        output = self.model(data)
        loss = self.loss_criterion(output, target)
        preds = torch.argmax(output, dim=1)
        acc = accuracy(preds, target)

        self.log(f"{stage}_loss", loss, prog_bar=True, on_epoch=True, on_step=False)
        self.log(f"{stage}_acc", acc, prog_bar=True, on_epoch=True, on_step=False)

        return preds, target #{'test loss': loss, 'test acc': acc, 'preds': preds, 'targets': target}

    def validation_step(self, batch, batch_idx):
        return self.evaluate(batch, "val")


    def validation_epoch_end(self, outputs):
        outs = tuple(map(torch.cat, zip(*outputs)))
        all_preds = outs[0]
        all_targets = outs[1]
        acc = accuracy(all_preds, all_targets)
        self.log('validation_acc_epoch', acc, on_epoch=True, on_step=False, prog_bar=True)
        return acc

    def test_step(self, batch, batch_idx):
        self.evaluate(batch, "test")

    def configure_optimizers(self):
        optimizer = optim.Adam(params=self.model.parameters(), lr=self.learning_rate)
        return optimizer

The following part is used for loading the data and training the model:

    use_cuda = True
    device = torch.device("cuda" if (use_cuda and torch.cuda.is_available()) else "cpu")

    train_data = torchvision.datasets.MNIST(data_dirs_dict[data_name], train=True, download=True, transform=transforms.ToTensor())
    test_data = torchvision.datasets.MNIST(data_dirs_dict[data_name], train=False, download=True, transform=transforms.ToTensor())

    train_loader = DataLoader(train_data,
                              batch_size=batch_size,
                              shuffle=True,
                              pin_memory=True)
    num_batches_train = len(train_loader)

    num_classes = len(train_data.classes)
    ## set up the dataloader for the optimizer steps:
    mpc_sampler = MPerClassSampler(train_data.targets, nopco, nopco * num_classes)

    mpc_dataloader = DataLoader(train_data,
                                batch_size=nopco * num_classes,
                                sampler=mpc_sampler,
                                pin_memory=True)

    test_loader = DataLoader(test_data,
                             batch_size=batch_size,
                             shuffle=False,
                             pin_memory=True)
                             
     ## Defining optimizers, loss and model:
     
    criterion = torch.nn.CrossEntropyLoss().to(device)
    optimizer1 = optim.Adam(params=net.parameters(), lr=lr)
    model = c_optimizer_lightning(net, learning_rate=lr, learning_rate_ratio=lrr,
                                                  loss_criterion=criterion, epoch_thresh=epoch_thresh,
                                                  custom_loss_function=uo.compute_custom_train_loss,
                                                  comparison_dataloader=mpc_dataloader, layers_id=layers_id,
                                                  matrix_norm_flag=matrix_norm_flag,
                                                  ).to(device)

 
    print(f'\nModel is on device: {model.device}\n')
    trainer = pl.Trainer(accelerator='auto', gpus=1, max_epochs=epochs)
    trainer.fit(model, train_loader, test_loader)

[akihironitta]

@hensel-f Could you provide your complete script and error so that I see where exactly the error is being raised?

As I quickly skimmed in your partial script, I realised that you're calling .to(device) which is not necessary in PyTorch Lightning. For more details, see our doc here: https://pytorch-lightning.readthedocs.io/en/stable/starter/converting.html#remove-any-cuda-or-to-device-calls

[jko55]

Hello, I got the same error when trying to reproduce "Lightning in 15 minutes" example from documentation

Adding the code at yhe end of tuturial

#embed 4 fake images!
fake_image_batch = Tensor(4, 28 * 28)
embeddings = encoder(fake_image_batch)
print("⚡" * 20, "\nPredictions (4 image embeddings):\n", embeddings, "\n", "⚡" * 20)

The line "embeddings = encoder(fake_image_batch)" throw this error

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
(when checking argument for argument mat1 in method wrapper_CUDA_addmm)

# linear.py  line 114
#113 def forward(self, input: Tensor) -> Tensor:
#114    return F.linear(input, self.weight, self.bias)     <<<<

accelerator = "gpu", devices="1"
or accelerator = "cpu" don't resolve the problem.

As Lighning take care for us of using cpu or gpu, i think it's a Lightning issue.

I'm using lightning-2.0.2 + pytorch 2.0.1+cu118 on windows 10

[Engineero]

I'm seeing a similar error and would love to know the resolution.

[jaroslawjanas]

I would love to see better documentation on the Lightning website.
Especially when it comes to issues like these.
I think right not it is lacking a bit and I can't find any information or even a starting point to begin diagnosing this issue.