import torch


def bucketize(tensor, bucket_boundaries):
    """Equivalent to numpy.digitize

    Notes
    -----
    Torch does not have a built in equivalent yet. I found this snippet here:
    https://github.com/pytorch/pytorch/issues/7284
    """
    result = torch.zeros_like(tensor, dtype=torch.int32)
    for boundary in bucket_boundaries:
        result += (tensor > boundary).int()
    return result