def METHOD_NAME(self, test_data, device, args):
    logging.info(f"----------test_on_the_client {args.rank} @ round {args.round_idx}--------")
    model = self.model
    model.to(device)
    model.eval()
    metrics = {"test_correct": 0, "test_loss": 0, "test_total": 0}
    criterion = torch.nn.CrossEntropyLoss().to(device)
    with torch.no_grad():
        for batch_idx, batch in enumerate(test_data):
            if args.model_class == "transformer":
                x = batch[1].to(device)
                target = batch[4].to(device)
            else:
                x, target = batch[0].to(device), batch[1].to(device)
            # x = x.to(device)
            # target = target.to(device)
            pred = model(x)
            if args.model_class == "transformer":
                pred = pred[0]
            loss = criterion(pred, target)
            _, predicted = torch.max(pred, -1)
            correct = predicted.eq(target).sum()
            metrics["test_correct"] += correct.item()
            metrics["test_loss"] += loss.item() * target.size(0)
            metrics["test_total"] += target.size(0)
    return metrics