Combine multiple autograd::Function or nest them?

Viewed 22

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?

0 Answers
Related