import torch


def gather_nodes(x, index):
    """
    Args:
        x:      (B, N, F) or (B, N)
        index:  (B, m)
    Returns:
        (B, m, F)
    """
    B, m = index.size()
    if len(x.size()) == 3:
        _, N, F = x.size()
        index = index.unsqueeze(-1).expand(B, m, F)
    elif len(x.size()) == 2:
        pass
    else:
        assert False, 'The size of `x` must either be (B, N, F) or (B, N).'

    return torch.gather(x, dim=1, index=index)