I think it will be more convenient this way 
def train(dataloader):
    model.train()
    
    total_acc, total_count = 0, 0
    total_loss = 0
    for idx, (label, text) in enumerate(dataloader):
        optimizer.zero_grad()
        predicted = model(text)
        loss = criterion(predicted, label)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        total_acc += (torch.round(torch.sigmoid(predicted)) == label).sum().item()
        total_count += label.size(0)
    return total_acc, total_loss, total_count