Thank you very much for your response.
So, as long as I understood, in your solution the function g(.) is assumed to be fixed; because you have returned None as its gradient. How can I also modify the gradient of g(.) assuming that it has some trainable parameters? (what needs to be returned instead of None, since g(.) is not a Tensor, but a module.)
Thanks again.