I am trying to export the model to onnx.
I can achieve that using the torch.onnx.export()
class net_nn(torch.nn.Module):
   def __init__(self, training=True):
       ....
   def load_weights(self, path, args=None):
       ....
   def foward(self, x):
       ....

model = net_nn(training=False)
model .load_weights(args.trained_model, args=args)

dummy_input = torch.ones(2, 3, dtype=torch.long)

torch.onnx.export(model, dummy_input , 'net_nn.onnx', verbose=True)

But what I want to do further is adding the post-processing in the onnx. Here is what I did,
class net_nn(torch.nn.Module):
   def __init__(self, training=True):
       ....
   def load_weights(self, path, args=None):
       ....
   def foward(self, x):
       ....

model = net_nn(training=False)
model .load_weights(args.trained_model, args=args)


def post_processing(x):
    ....

class net_and_post(torch.nn.Module):
      def foward(self, x):
        pred = model(x)
        return post_processing(pred)

model_with_post = net_and_post()
dummy_input = torch.ones(2, 3, dtype=torch.long)
torch.onnx.export(model_with_post , dummy_input , 'net_nn_post.onnx', verbose=True)

But it ends by
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
Please help.