def METHOD_NAME(self, batch, g, h, c):
    """Compute tree-lstm prediction given a batch.
    Parameters
    ----------
    batch : dgl.data.SSTBatch
        The data batch.
    g : dgl.DGLGraph
        Tree for computation.
    h : Tensor
        Initial hidden state.
    c : Tensor
        Initial cell state.
    Returns
    -------
    logits : Tensor
        The prediction of each node.
    """
    # feed embedding
    embeds = self.embedding(batch.wordid * batch.mask)
    g.ndata["iou"] = self.cell.W_iou(
        self.dropout(embeds)
    ) * batch.mask.float().unsqueeze(-1)
    g.ndata["h"] = h
    g.ndata["c"] = c
    # propagate
    dgl.prop_nodes_topo(
        g,
        self.cell.message_func,
        self.cell.reduce_func,
        apply_node_func=self.cell.apply_node_func,
    )
    # compute logits
    h = self.dropout(g.ndata.pop("h"))
    logits = self.linear(h)
    return logits