BUMP!
Somewhat related: I am trying to save my model as shown below:

    with torch.no_grad():
        input_shape = [1, 3, 512, 256]
        input_data = torch.randn(input_shape)
        input_data = input_data.to(DEVICE) # GPU
        out_res = model(input_data)
        scripted_model = torch.jit.trace(model, input_data)
        scripted_model.eval()
    # Save your model. The following code saves it with the .pth file extension
    scripted_model.save('model.pth')

I keep getting this error which makes no sense to me (since I am adding the no_grad() )
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
Tensor:
(1,1,.,.) =
 0.01 * ...

Any hints would be very helpful. Thanks.