def cholesky(u):
    """
    Like :func:`torch.cholesky` but uses sqrt for scalar matrices.
    Works around https://github.com/pytorch/pytorch/issues/24403 often.
    """
    if u.size(-1) == 1:
        return u.sqrt()
    return u.cholesky()
