def METHOD_NAME(self, train_data, device, args):
    model = self.model
    model.to(device)
    model.METHOD_NAME()
    test_data = None
    try:
        test_data = self.test_data
    except:
        pass
    criterion = torch.nn.MSELoss() if args.dataset != "qm9" else torch.nn.MAELoss()
    if args.client_optimizer == "sgd":
        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate)
    else:
        optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
    min_score = np.Inf if args.metric != "r2" else -np.Inf
    best_model_params = {}
    # print('Training on {}'.format(torch.cuda.get_device_name()))
    for epoch in range(args.epochs):
        avg_loss = 0
        count = 0
        for mol_idxs, (adj_matrix, feature_matrix, label, _) in enumerate(
            train_data
        ):
            optimizer.zero_grad()
            adj_matrix = adj_matrix.to(
                device=device, dtype=torch.float32, non_blocking=True
            )
            feature_matrix = feature_matrix.to(
                device=device, dtype=torch.float32, non_blocking=True
            )
            label = label.to(device=device, dtype=torch.float32, non_blocking=True)
            logits = model(adj_matrix, feature_matrix)
            loss = criterion(logits, label)
            loss.backward()
            optimizer.step()
            if test_data is not None:
                test_score, _ = self.test(self.test_data, device, args)
                if args.metric != "r2":
                    print(
                        "Epoch = {}: Test {} = {}".format(
                            epoch, args.metric.upper(), test_score
                        )
                    )
                    if test_score < min_score:
                        min_score = test_score
                        best_model_params = {
                            k: v.cpu() for k, v in model.state_dict().items()
                        }
                    print(
                        "Current best {}= {}".format(args.metric.upper(), min_score)
                    )
                else:
                    print("Epoch = {}: Test R2 = {}".format(epoch, test_score))
                    if test_score > min_score:
                        min_score = test_score
                        best_model_params = {
                            k: v.cpu() for k, v in model.state_dict().items()
                        }
                    print("Current best R2= {}".format(min_score))
        #
        #     avg_loss += loss.item()
        #     count += 1
        #     # logging.info("training. epoch = %d, mol_idxs = %d, avg_loss = %f" % (epoch, mol_idxs, avg_loss))
        #
        # avg_loss /= count
    return min_score, best_model_params