Hi @eqy, related to this, do you know how I can use this approach of iterating class by class so that I can compute the accuracy for each class?
This is what I currently have. check_accuracy gets executed for each epoch during training.
def check_accuracy(loader, model, device="cuda"):
    model.eval()

    with torch.no_grad():
        for x, y, _ in loader:
            x = x.to(device)
            y = y.to(device).unsqueeze(1)
            preds = model(x)
            num_classes = preds.shape[1]
            for class_idx in range(num_classes):
                curr_class = (preds[:,class_idx,:,:] == torch.max(preds, dim=1)[0]).float()
                curr_class = curr_class.unsqueeze(1)
                num_correct = (curr_class == y).float()
                class_accuracy = num_correct.sum()/num_correct.numel()
                class_accuracy = class_accuracy*100
                print(f"{class_accuracy}% accuracy for class {class_idx}\n")
    model.train()

Is this the correct way to do it?
Thank you very much!