import torch


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