Injecting gradient directly into the training extension

[berndporr]

I do realtime learning with libtorch where the gradient is calculated elsewhere and would like to translate that to executorch: https://github.com/berndporr/dnf_torch this is realtime noise cancellations but I have also other applications where the gradient comes further afield so really not associated with a simple subtraction operation as I also do RL. In libtorch I can directly inject the gradient and then do the backwards pass: https://github.com/berndporr/dnf_torch/blob/main/dnf_torch.cpp#L99

torch::Tensor gradient = torch::tensor({-f_nn}).to(device);
output.backward(gradient);
optimizer.step();

However, in executorch forward/backward is in one go. How do I go about this?
Looking at the XOR example I could inject the gradient as an additional input? Roughly:

class TrainingNet(nn.Module):
    def __init__(self, net):
        super().__init__()
        self.net = net

    def forward(self, gradient, input):
        pred = self.net(input)
        return gradient, pred.detach().argmax(dim=1)

and then lower that and inject the gradient as the first input. Would that be the way to go?