def METHOD_NAME(self, fe_index, s):
    """
    fe_index: full_edge_index, [2, all_edges_batchwise]
    s: predicted edge value, [all_edges_batchwise, 1]
    construct the predicted edge set and corresponding edge weights
    """
    s = s[:, 0]
    fe_index = paddle.transpose(fe_index, perm=[1, 0])
    sender = paddle.unsqueeze(fe_index[0][s > 0], 0)
    receiver = paddle.unsqueeze(fe_index[1][s > 0], 0)
    pred_index = paddle.concat([sender, receiver], 0)
    pred_weight = s[s > 0]
    pred_index = paddle.transpose(pred_index, perm=[1, 0])
    return pred_index, pred_weight