Error while trying to log accuracy

[sushmit86]

Hi all -
I am new to using Pytorch as well as pytorch_lightning. I am getting an error while trying to log accuracy during the training step.
Here is my code below

import torch
from torch import nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader, random_split
from torch.nn import functional as F
from torchvision.datasets import MNIST,FashionMNIST
from torchvision import datasets, transforms
import os


class LightningMNISTClassifier(pl.LightningModule):

  def __init__(self):
    super().__init__()
    self.model = torch.nn.Sequential(nn.Conv2d(1, 6, kernel_size=5,
                                               padding=2), nn.Sigmoid(),
                          nn.AvgPool2d(kernel_size=2, stride=2),
                          nn.Conv2d(6, 16, kernel_size=5), nn.Sigmoid(),
                          nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(),
                          nn.Linear(16 * 5 * 5, 120), nn.Sigmoid(),
                          nn.Linear(120, 84), nn.Sigmoid(), nn.Linear(84, 10))

  def forward(self, x):
      x = x.view(-1, 1, 28, 28)
      x = self.model(x)
      return x

  def loss_fn(self, logits, y):
    loss = nn.CrossEntropyLoss()
    return loss(logits,y)

  def training_step(self, train_batch, batch_idx):
      x, y = train_batch
      logits = self.forward(x)
      loss = self.loss_fn(logits, y)
      self.log('train_loss', loss)
      m = nn.Softmax(dim=1)
      output = m(logits)
      accuracy = Accuracy(top_k=1)
      self.log('train_accuracy', accuracy(output,y))
      return loss


  def validation_step(self, val_batch, batch_idx):
      x, y = val_batch
      logits = self.forward(x)
      loss = self.loss_fn(logits, y)
      self.log('val_loss', loss)

  def configure_optimizers(self):
    optimizer = torch.optim.SGD(self.model.parameters(), lr=0.9)
    return optimizer



class MNISTDataModule(pl.LightningDataModule):

  def setup(self, stage):
    # transforms for images
    transform=transforms.Compose([transforms.ToTensor()])
      
    # prepare transforms standard to MNIST
    self.mnist_train = FashionMNIST(root="../data", train=True, download=True, transform=transform)
    self.mnist_test = FashionMNIST(root="../data", train=False, download=True, transform=transform)

  def train_dataloader(self):
    return DataLoader(self.mnist_train, batch_size=256,shuffle=True,num_workers=1)

  def val_dataloader(self):
    return DataLoader(self.mnist_test, batch_size=256,num_workers=1)


data_module = MNISTDataModule()
# train
model  = LightningMNISTClassifier()
trainer = pl.Trainer(max_epochs=1,gpus=1,progress_bar_refresh_rate =1)
trainer.fit(model, data_module)

However I am getting an error

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Any help on this

[awaelchli]

Hi, should be an easy fix. The problem is in training step where you define

accuracy = Accuracy(top_k=1)

but it should be defined in the constructor

self.accuracy = Accuracy(top_k=1)
and in training_step you call

self.accuracy(output, y)

Hope this helps.

[sushmit86]

Thanks for the reply