I have a process that looks something like this:
module_one = #.... some NN
module_two = #.... some NN

X2 = module_one(X1)
Y = module_two(X2)

I want to manually get the cost for X2 as resulting from the backwards pass of model one and pass it as the cost to module_two
So something like
X2.required_grad = True
Y = module_two(X2)
cost = X2.grad.cpu().detach()

module_one.attach_cost(cost)
cost.bakcwards()
module_one_optim.step()

I am unsure how to accomplish the attach_cost step.
I know that for this simple example I could just combine the forward pass for both modules into one… in practice, due to the structure of my code this is rather hard, so I’d prefer to just manually attach the cost and do the backprop based on that.
Is this possible?