def train(model, optimizer, criterion, x, y_true, train_idx):
    """
        One step of optimization for DNN

            Arguments:
                model : DNN
                    DNN model
                optimizer: Optimizer
                    Optimization function
                criterion: 
                    Loss function
                x : torch.tensor of shape (num_nodes, num_node_features)
                    Node feature matrix
                y_true : torch.tensor of shape (num_nodes)
                    True node labels
                train_idx : torch.tensor of shape (num_nodes)
                    Training node index
                
            Returns:
                loss : int
                    Training loss
    """
    model.train()
    
    optimizer.zero_grad()
    _, out, _ = model(x[train_idx])
    loss = criterion(out, y_true.squeeze(1)[train_idx])
    loss.backward()
    optimizer.step()

    return loss.item()