My bad. It was supposed to be fmodel in the original change itself. However, even after using fmodel, you’ll notice that the torch.autograd.grad return all zeros. And this I think is because the backward graph doesn’t extend to model A. Trying to find out how to do this so that the grads are not zero!