Saving a Pytorch neural net (torch.autograd.grad included) as a Torch Script code

Question

I have the following neural net setup:

enter image description here

The neural net takes values of x as input, and is trained to output f(x) = x**2. The neural net is also trained to output the first-order derivative of f(x), i.e. f’(x) = 2*x. This derivative is computed using torch.autograd.grad, and is included as a second term in the loss function.

I have been unable to compile the neural net as a TorchScript code - for the example above, this is the error I have:

RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient Tensor: 0.2816 -0.1232 -0.1976 0.3143 [ torch.FloatTensor{4,1} ]

I have tried readjusting the neural net multiples times, but an error always comes up, and it always has to do with using torch.autograd.grad. How do I go about saving the above neural net as a Torchscript code? What am I missing here?

Thanks!