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 ?