I think found the solution myself (through trial and error). The argument “create_graph=True” should be added to torch.autograd.grad().
I am still open to discuss this though, if there is any error in it