def METHOD_NAME(self, cuda=False):
    device = torch.device("cuda") if cuda else torch.device("cpu")
    for dtype in (torch.float, torch.double):
        for batch_mode in (False, True):
            if batch_mode:
                A = self._gen_test_psd().to(device=device, dtype=dtype)
                D = torch.eye(2).type_as(A).unsqueeze(0).repeat(2, 1, 1)
            else:
                A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
                D = torch.eye(2).type_as(A)
            A += D
            # basic
            L = torch.linalg.cholesky(A)
            L_safe = psd_safe_cholesky(A)
            self.assertTrue(torch.allclose(L, L_safe))
            # upper
            L = torch.linalg.cholesky(A).mT
            L_safe = psd_safe_cholesky(A, upper=True)
            self.assertTrue(torch.allclose(L, L_safe))
            # output tensors
            L = torch.empty_like(A)
            L_safe = torch.empty_like(A)
            torch.linalg.cholesky(A, out=L)
            psd_safe_cholesky(A, out=L_safe)
            self.assertTrue(torch.allclose(L, L_safe))
            # output tensors, upper
            torch.linalg.cholesky(A, out=L)
            L = L.mT
            psd_safe_cholesky(A, upper=True, out=L_safe)
            self.assertTrue(torch.allclose(L, L_safe))
            # make sure jitter doesn't do anything if p.d.
            L = torch.linalg.cholesky(A)
            L_safe = psd_safe_cholesky(A, jitter=1e-2)
            self.assertTrue(torch.allclose(L, L_safe))