def train(inputs, weight1, weight2, adj_matrix, am_partitions, optimizer, local_train_mask, local_labels, device):
    outputs = GCNFunc.apply(inputs, weight1, adj_matrix, am_partitions, F.relu, 'layer1')
    outputs = GCNFunc.apply(outputs, weight2, adj_matrix, am_partitions, (lambda x: F.log_softmax(x, dim=1)), 'layer2')
    optimizer.zero_grad()
    if (list(local_labels[local_train_mask].size())[0] > 0):
        loss = F.nll_loss(outputs[local_train_mask], local_labels[local_train_mask])
        loss.backward()
    else:
        fake_loss = (outputs * torch.cuda.FloatTensor(outputs.size(), device=device).fill_(0)).sum()
        fake_loss.backward()
    optimizer.step()
    return outputs
def train(inputs, weight1, weight2, adj_matrix, am_partitions, optimizer, local_train_mask, local_labels, device):
    outputs = GCNFunc.apply(inputs, weight1, adj_matrix, am_partitions, F.relu, 'layer1')
    outputs = GCNFunc.apply(outputs, weight2, adj_matrix, am_partitions, (lambda x: F.log_softmax(x, dim=1)), 'layer2')
    optimizer.zero_grad()
    if (list(local_labels[local_train_mask].size())[0] > 0):
        loss = F.nll_loss(outputs[local_train_mask], local_labels[local_train_mask])
        loss.backward()
    else:
        fake_loss = (outputs * torch.cuda.FloatTensor(outputs.size(), device=device).fill_(0)).sum()
        fake_loss.backward()
    optimizer.step()
    return outputs