diff --git a/test/test_ops.py b/test/test_ops.py
index 51f9f0f..4e72190 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -30,6 +30,7 @@
   "_segment_reduce",
   "bincount",  # NOTE: dtype for int input torch gives float. This is weird.
   "byte",
+  "cholesky",  # To be removed from PyTorch, still covered by linalg.cholesky
   "cholesky_solve",
   "geqrf",
   "histogram",  # hard op: AssertionError: Tensor-likes are not close!
diff --git a/torchax/amp.py b/torchax/amp.py
index ae5d5cf..2007dba 100644
--- a/torchax/amp.py
+++ b/torchax/amp.py
@@ -128,7 +128,6 @@ def autocast(device, dtype=torch.bfloat16, env=None):
   torch.ops.aten.grid_sampler_3d.default: CastPolicy.FP32,
   torch.ops.aten.trace.default: CastPolicy.FP32,
   torch.ops.aten.view_as_complex.default: CastPolicy.FP32,
-  torch.ops.aten.cholesky.default: CastPolicy.FP32,
   torch.ops.aten.cholesky_inverse.default: CastPolicy.FP32,
   torch.ops.aten.cholesky_solve.default: CastPolicy.FP32,
   torch.ops.aten.inverse.default: CastPolicy.FP32,
diff --git a/torchax/ops/jaten.py b/torchax/ops/jaten.py
index e2d72d2..9bd3643 100644
--- a/torchax/ops/jaten.py
+++ b/torchax/ops/jaten.py
@@ -203,7 +203,7 @@ def _aten_index_select(x, dim, index):
   return jnp.take(x, index, dim)
 
 
-@op(torch.ops.aten.cholesky)
+@op(torch.ops.aten.linalg_cholesky)
 def _aten_cholesky(input, upper=False):
   return jax.scipy.linalg.cholesky(input, lower=(not upper))