# Error when implementing custom autograd.Function

**URL:** <https://discuss.pytorch.org/t/error-when-implementing-custom-autograd-function/212031>\
**Category:** autograd\
**Created:** [October 24, 2024, 11:27am UTC](https://discuss.pytorch.org/t/error-when-implementing-custom-autograd-function/212031 "2024-10-24T11:27:10Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Elling\_Svee](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/elling_svee/32/73442_2.png) [@Elling\_Svee](https://discuss.pytorch.org/u/Elling_Svee)\
**Post date:** [October 24, 2024, 11:27am UTC](https://discuss.pytorch.org/t/error-when-implementing-custom-autograd-function/212031/1 "2024-10-24T11:27:10Z")

</div>

I am currently learning how to implement custom `autograd.Functions`, and for the sake of practice I try to re-write the backwards pass of the `torch.linalg.cholesky`-function. The code in the implemented backwards pass is intended to be the same as the `cholesky_backward` function at line 1939 in [FunctionsManual.cpp](https://github.com/pytorch/pytorch/blob/96b30dcb25c80513769dae2a8688aec080b00117/torch/csrc/autograd/FunctionsManual.cpp). The only difference is that i have written it in Python, and attempted to avoid in-place operations.

However, the code does not pass the `gradcheck`, which raises as error stating that there is a “Jacobian mismatch”. Can anyone tell me what I am doing wrong?

```auto
import torch
class CustomCholFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A: torch.Tensor) -> torch.Tensor:
        A_cpu = A.cpu().double()

        # Perform Cholesky decomposition
        L = torch.linalg.cholesky(A_cpu, upper=False)

        # Save the lower triangular matrix for backward pass
        L_cpu = L.cpu()
        ctx.L_cpu = L_cpu

        return L_cpu.to(A.device, dtype=A.dtype)

    @staticmethod
    def backward(ctx, gL: torch.Tensor) -> torch.Tensor:
        gL_cpu = gL.cpu()
        L_cpu = ctx.L_cpu.clone()

        # This code is meant to be the same as the backwards pass in the
        # autograd implementation of torch.linalg.cholesky
        upper = False # Currently assume lower triangular
        L_ = L_cpu.adjoint() if upper else L_cpu
        gL_ = gL_cpu.adjoint() if upper else gL_cpu

        # Compute gradient of A using out-of-place operations
        gA = torch.matmul(L_.adjoint(), gL_).tril()

        gA = 0.5 * (gA + gA.tril(-1).adjoint())
        gA = torch.linalg.solve_triangular(L_.adjoint(), gA, upper=True, left=True)
        gA = torch.linalg.solve_triangular(L_, gA, upper=False, left=False)

        return gA.to(gL.device)

# Make SPD matrix
A = torch.randn(5, 5, dtype=torch.double)
A = A @ A.t() + 0.1 * torch.eye(5, dtype=torch.double)
A.requires_grad = True

# Works as intended
CustomCholFunc.apply(A)

# Raises an error
torch.autograd.gradcheck(CustomCholFunc.apply, A)

```