Using a custom PyTorch backward pass with Lightning

[thibmonsel]

Hello,

I'm currently using a custom backward with torch.autograd.Function and I can't seem to find anything that shows how to integrate it with Lightning. For a custom backward pass we could for example think of the one from Neural ODE's (Algorithm 1).

What I have is something of the sort :

class integrate(torch.autograd.Function):
    @staticmethod
    def forward(ctx, args, *params):
        ....
        return y

   @staticmethod
    def backward(ctx, *grad_y):
        ...
        return None, *grad_params

For a regular Pytorch training I would call it as something of the following :

for i in range(max_epoch):
    opt.zero_grad()
    out = integrate(args, model, ts)
    loss = lossfunc(out, y_truth)
    loss.backward()
    opt.step()

How would I integrate this in a LightningModule ? As far as I know, the training_step and backward methods won't do the trick ? Maybe I'm not thinking of other hooks !

Another possibility would be to use Legendre example.
Thanks in advance !