how to export the module

when i export the UPerNet like this:
traced_script_module = torch.jit.trace(segmentation_module, feed_dict['img_data'])
traced_script_module.save("seg.pt")
it does not work.
then i try export the encoder and decoder seperately, like this:
encoder = segmentation_module.get_encoder()
jit_encoder = torch.jit.trace(encoder, feed_dict['img_data'])
jit_encoder.save('encoder.pt')
but this assertion failed：
assert(isinstance(orig, torch.nn.Module))
AssertionError

any idea？thanks in advance