import torch


def get_mask(
    ref, 
    sub, 
    pyg=False
):
    """
    Get the mask for a reference list based on a subset list.

    Args:
        ref: reference list
        sub: subset list
        pyg: boolean; whether to return torch tensor for PyG

    Return:
        mask: list or torch.BoolTensor
    """
    mask = [item in sub for item in ref]
    if pyg:
        mask = torch.BoolTensor(mask)
    return mask