When developing a custom autograd::Function, I can write something like this:
struct my_func: torch::autograd::Function<my_func>{
static torch::Tensor forward(torch::autograd::AutogradContext *ctx, //param list)
{
// do forward computation
ctx->save_for_backward({//some list of tensors});
return a_tensor;
}
static torch::autograd::tensor_list backward(torch::autograd::AutogradContext *ctx, torch::autograd::tensor_list grad_out){
auto saved = ctx->get_saved_variables();
// do backward computation
return {grad_tensor list};
}
};
struct GNNLayer : torch::nn::Module{
torch::Tensor forward(//paramlist) {
torch::Tensor x = my_func::apply(//arg list);
return layer_output;
}
};
What I want to do now is write multiple such autograd::Function and invoke them inside my nn::Module while passing the gradients between them. What is the syntax to do that?