import torch


def get_default_dropout(p=0., dim=2, inplace=True):
    """Wrapper for torch's dropout.
    """
    if dim == 1:
        return torch.nn.Dropout(p=p, inplace=inplace)
    elif dim == 2:
        return torch.nn.Dropout2d(p=p, inplace=inplace)
    elif dim == 3:
        return torch.nn.Dropout3d(p=p, inplace=inplace)
    else:
        raise TypeError("Expected dim = 1, 2 or 3, but got {}".format(dim))