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)
            else:
                A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
            idx = torch.arange(A.shape[-1], device=A.device)
            # default values
            Aprime = A.clone()
            Aprime[..., idx, idx] += 1e-6 if A.dtype == torch.float32 else 1e-8
            L_exp = torch.linalg.cholesky(Aprime)
            with warnings.catch_warnings(record=True) as w:
                # Makes sure warnings we catch don't cause `-w error` to fail
                warnings.simplefilter("always", NumericalWarning)
                L_safe = psd_safe_cholesky(A)
                self.assertTrue(any(issubclass(w_.category, NumericalWarning) for w_ in w))
                self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
            self.assertTrue(torch.allclose(L_exp, L_safe))
            # user-defined value
            Aprime = A.clone()
            Aprime[..., idx, idx] += 1e-2
            L_exp = torch.linalg.cholesky(Aprime)
            with warnings.catch_warnings(record=True) as w:
                # Makes sure warnings we catch don't cause `-w error` to fail
                warnings.simplefilter("always", NumericalWarning)
                L_safe = psd_safe_cholesky(A, jitter=1e-2)
                self.assertTrue(any(issubclass(w_.category, NumericalWarning) for w_ in w))
                self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
            self.assertTrue(torch.allclose(L_exp, L_safe))