You would need to implement the backward method manually for this custom autograd.Function, as Autograd doesn’t know how the gradients should be passed through this layer.
This can be avoided if you could use pure PyTorch methods in the forward pass, but this doesn’t seem to be the case, as you are apparently using methods from other libraries.
This tutorial 46 gives you an example of a custom function.