def METHOD_NAME(self,
            node_features,
            neighbor_features,
            edge_features,
            edge_index):
    """Apply NNConv layer.
    Args:
        node_features (Layer): A 2D layer of node features of
                               shape (num_nodes, input_channels)
        neighbor_features (Layer): A 3D layer of node features of
                                   shape (num_edges, 1, input_channels)
        edge_features (Layer): A 2D layer of edge features of
                               shape (num_edges, edge_features)
        edge_index (Layer): A 1D layer of node features of
                            shape (num_edges * output_channels).
                            The indices used for reduction
    Returns:
        (Layer): The output after NNConv. The output layer has the shape
                 (num_nodes, self.output_channels)
    """
    updated_node_fts, neighbor_vals = self.message(node_features,
                                                   neighbor_features,
                                                   edge_features)
    aggregated_fts = self.aggregate(neighbor_vals, edge_index)
    update = lbann.Sum(updated_node_fts,
                       aggregated_fts,
                       name=self.name+"_updated_node_features")
    return update