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.BCEWithLogitsLoss(reduction="none")
    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)
    max_test_score = 0
    best_model_params = {}
    for epoch in range(args.epochs):
        for mol_idxs, (forest, feature_matrix, label, mask) in enumerate(
            train_data
        ):
            # Pass on molecules that have no labels
            if torch.all(mask == 0).item():
                continue
            optimizer.zero_grad()
            forest = [level.to(device=device, dtype=torch.long, non_blocking=True) for level in forest]
            feature_matrix = feature_matrix.to(
                device=device, dtype=torch.float32, non_blocking=True
            )
            label = label.to(device=device, dtype=torch.float32, non_blocking=True)
            mask = mask.to(device=device, dtype=torch.float32, non_blocking=True)
            logits = model(forest, feature_matrix)
            loss = criterion(logits, label) * mask
            loss = loss.sum() / mask.sum()
            loss.backward()
            optimizer.step()
            if ((mol_idxs + 1) % args.frequency_of_the_test == 0) or (
                mol_idxs == len(train_data) - 1
            ):
                if test_data is not None:
                    test_score, _ = self.test(self.test_data, device, args)
                    print(
                        "Epoch = {}, Iter = {}/{}: Test Score = {}".format(
                            epoch, mol_idxs + 1, len(train_data), test_score
                        )
                    )
                    if test_score > max_test_score:
                        max_test_score = test_score
                        best_model_params = {
                            k: v.cpu() for k, v in model.state_dict().items()
                        }
                    print("Current best = {}".format(max_test_score))
    return max_test_score, best_model_params