Problem calculating accuracy on multiclass classfication

Hi, I am facing a problem when I try to compute accuracy on a multiclass classification problem with the Accuracy() function from torchmetrics
Here is the training function I created:

from torchmetrics import Accuracy
accuracy = Accuracy()

def training(epoch, model, train_loader, optimizer, criterion):
  "Training over an epoch"
  metric_monitor = MetricMonitor()
  
  model.train()

  for batch in train_loader:
    images = batch['t1'][tio.DATA].cuda()
    labels = batch['label'].cuda()
    output = F.softmax(model(images), dim=0)

    loss = criterion(output, labels)

    output = output.data.max(dim=1,keepdim=True)[1]

    acc = accuracy(output, labels)
 
    metric_monitor.update("Loss", loss.item())
    metric_monitor.update("Accuracy", acc)

    optimizer.zero_grad()

    loss.backward()
    optimizer.step()
  print("[Epoch: {epoch:03d}] Train      | {metric_monitor}".format(epoch=epoch, metric_monitor=metric_monitor))
  return metric_monitor.metrics['Loss']['avg'], metric_monitor.metrics['Accuracy']['avg']

When I try to call this function to train the model I get the following error :
ValueError: If preds have one dimension more than target, preds should be a float tensor.
Does anybody know what I am doing wrong ?

@gkrisp9 This is a classic case of CrossEntropy evaluation

https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html

When you have multi label classification, your output dimension becomes
batch_size, num_of_classes

And your input target is batch_size

The solution would be to add this extra line

output = output.data.max(dim=1,keepdim=True)[1]
output = output.view(-1)
acc = accuracy(output, labels)

Thanks again! Now I am getting this error: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
I am bit confused because as I see from my code all the labels etc are in cuda.
I added this command: output.cuda() to see if that was the problem but the error appears again.

Update: I made a function to return accuracy and pass output and labels to cpu and it works

output = F.softmax(model(images), dim=0).cuda()
...
output = output.data.max(dim=1,keepdim=True)[1].cuda()