import torch


def avg_accuracy(prob_vecs, labels):
  """Example metrics function: average accuracy
  Let N := number of data points and C := the number of classes
  
  Args:
    prob_vecs (torch.Tensor): Tensor of size (N, C)
    labels (torch.Tensor): Tensor of size (N), hard labels, with classes corresponding to indices 0, ..., C-1
  
  Returns:
    torch.Tensor: Tensor of size (N), average accuracy over all datapoints.
  """
  hard_preds = torch.argmax(prob_vecs, 1)
  return (hard_preds == labels).float().mean()