Solved by ptrblck in post #5 
              
              
                The batch_size in my code example would correspond to the batch size you set in your DataLoader. 
To update the conf mat you would have to pass and return it from the method: 
def confusion_matrix(preds, labels, conf_matrix):
    preds = torch.argmax(preds, 1)
    for p, t in zip(preds, labels):
   …