I write a forward function by using CUDAExtension kernal in device cuda. I want to use the autograd model to compute the gradient of this processing. Which document can help me to complete this work?
Assuming you are using your custom extension in an autograd.Function you would need to implement the backward method manually as well. PyTorch will handle the backward automatically if you use native and differentiable PyTorch ops.
Thank you for your help, Sir!
But I am confused in this part. If I write the backward function by CUDAExtension kernal,where can autograd be demonstrated? Are the backward function and the gradient equivalent？Perhaps there is something wrong with my understanding.Does this mean that if I use the custom cuda extension function, I cannot take advantage of automatic differentiation?
The backward function computes the gradients for all inputs as seen in e.g. this tutorial.
Yes, as PyTorch won’t be able to reason about the custom operations you are using unless you use PyTorch operations.
Thank you for you help, Sir! Your remark answer my doubts.