def get_batch_size(tensor, base_size):
    """Get the batch size of a tensor if it is a discrete or continuous tensor.

    Parameters
    ----------
    tensor: torch.tensor.
        Tensor to identify batch size
    base_size: Tuple.
        Base size of tensor.

    Returns
    -------
    batch_size: int or None

    """
    size = tensor.shape
    if len(base_size) == 0:  # Discrete
        return tuple(size)
    else:
        return tuple(size[: -len(base_size)])