import torch


def _pad_norm(x, implicit=False):
    """Add a channel that ensures that prob sum to one if the input has
    an implicit background class. Else, ensures that prob sum to one."""
    if implicit:
        x = torch.cat((1 - x.sum(dim=1, keepdim=True), x), dim=1)
    return x