'auto' grad with 'custom' grad function

I have a torch model which first pass the torch Conv2d module and then pass my custom transformer encoder block(made into torch module).
And I made a custom grad function for transformer encoder block.
This transformer encoder block works well both forward and backward.
What I want is that, for transformer encoder, it uses my custom grad function and for conv net, it uses auto grad.
I tried train my model but it seems conv params are not updating.
Is there any way to use both autograd and custom grad?

Could you explain your custom autograd approach a bit more?
Did you write a custom autograd.Function or changed anything else?

Yes, I used autograd.Function

In that case Autograd will use your custom Function and I’m not sure how it would break the conv net.
Could you add more information about your use case and post a minimal, executable code snippet to reproduce the issue, please?