So I have to train two models simultaneously, where the input of the 2nd model is the output of the 1st model.
And the output of my 2nd model is the grad of 1st model’s output wrt its input
Previously I was training these 2 models independently

getting the outputs from the 1st model, then calculating its grad wrt input,
then feeding these outputs of 1st model to my 2nd model, and checking if the produced grad from 2nd model is equal to the original grads computed at the 1st step, using MSE loss.

Now I want to train all this simultaneously. And I cant think of how
I was calculating the grads like this
input.requires_grad = True 
opx = model(input)  
opx.backward(torch.ones_like(opx))
dx1 = input.grad.data

Help will be much appreciated.