I have a model that implements a layer with a custom forwards and backwards function following the extension guide 1. However the backwards function is not computed using the gradient of the downstream layer, but the gradient of the output layer.
What’s the cleanest way to pass it to the backwards function of upstream layers?