src/utils/metrics.py

import torch

def accuracy(logits, y):
    return (logits.argmax(1) == y).float().mean().item()

def topk_accuracy(logits, y, k=5):
    topk = logits.topk(k, dim=1).indices
    return (topk == y.unsqueeze(1)).any(dim=1).float().mean().item()