Error Tracing the model using torch.jit.trace

[satyajitghana]

While trying to export the model to traced pytorch model, i am getting Could not export Python function call 'Scatter' how do i resolve this ?

Notebook for reference: https://colab.research.google.com/drive/1vf8osp0aeQA7dgBTw5Ejy9OibGOvgpaj?usp=sharing

model = AnyNet(args)
model = nn.DataParallel(model).cuda()
checkpoint = torch.load('/content/checkpoint/kitti2015_ck/checkpoint.tar')
model.load_state_dict(checkpoint['state_dict'], strict=False)
example = [torch.rand(1, 3, 3840, 2160).cuda(), torch.rand(1, 3, 3840, 2160).cuda()]
traced_model = torch.jit.trace(model, example_inputs=example)
traced_model.save("anynet_3840x2160_b1.traced.pt")

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-15-462347c4b4a7> in <module>()
----> 1 traced_model.save("anynet_3840x2160_b1.traced.pt")

/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py in save(self, *args, **kwargs)
    485             See :func:`torch.jit.save <torch.jit.save>` for details.
    486             """
--> 487             return self._c.save(*args, **kwargs)
    488 
    489         def _save_for_lite_interpreter(self, *args, **kwargs):

RuntimeError: 
Could not export Python function call 'Scatter'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/scatter_gather.py(13): scatter_map
/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/scatter_gather.py(15): scatter_map
/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/scatter_gather.py(28): scatter
/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/scatter_gather.py(36): scatter_kwargs
/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/data_parallel.py(168): scatter
/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/data_parallel.py(157): forward
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py(709): _slow_forward
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py(725): _call_impl
/usr/local/lib/python3.7/dist-packages/torch/jit/_trace.py(940): trace_module
/usr/local/lib/python3.7/dist-packages/torch/jit/_trace.py(742): trace
<ipython-input-14-e92379b43790>(2): <module>

[satyajitghana]

fixed it

model = AnyNet(args)
model = nn.DataParallel(model).cuda()
checkpoint = torch.load('/content/checkpoint/kitti2015_ck/checkpoint.tar')
model.load_state_dict(checkpoint['state_dict'], strict=False)
torch.save(model.module.state_dict(), 'non-parallel-model.pth')

model_np = AnyNet(args)
model_np = model_np.cuda()
model_np.load_state_dict(torch.load('non-parallel-model.pth'), strict=False)
model_np = model_np.eval()

example = [torch.rand(1, 3, 1280, 720).cuda(), torch.rand(1, 3, 1280, 720).cuda()]
traced_model = torch.jit.trace(model_np, example_inputs=example)
traced_model.save("anynet_1280x720_b1.traced.pt")

the key was to save model.module and just loading them

[AloneGu]

try:

traced_model = torch.jit.trace(model, example_inputs=example)

to

traced_model = torch.jit.trace(model.module, example_inputs=example)

[shashankb2003]

When trying to use the following command:

traced_model = torch.jit.trace(model.module, example_inputs=example)

I am getting "Could not export Python function call 'GateRecurrent2dnoindFunction' ". How do I resolve this?
Is this error occurring because of custom forward functions? I tried changing the @staticmethod decorator to @script but to no avail

model = AnyNet(args)
model = nn.DataParallel(model).cuda()
checkpoint = torch.load('/content/checkpoint.tar')
model.load_state_dict(checkpoint['state_dict'], strict=False)
example = [torch.rand(1, 3, 1280, 720).cuda(), torch.rand(1, 3, 1280, 720).cuda()]
traced_model = torch.jit.trace(model.module, example_inputs=example)
traced_model.save("anynet_1280x720_b1.traced.pt")

RuntimeError: 
Could not export Python function call 'GateRecurrent2dnoindFunction'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py(539): apply
/content/AnyNet/models/spn_t1/modules/gaterecurrent2dnoind.py(12): forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1508): _slow_forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1527): _call_impl
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1518): _wrapped_call_impl
/content/AnyNet/models/anynet.py(166): forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1508): _slow_forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1527): _call_impl
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1518): _wrapped_call_impl
/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py(1065): trace_module
/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py(798): trace
<ipython-input-17-c0dc6b3a9cf8>(6): <cell line: 6>
/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py(3553): run_code
/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py(3473): run_ast_nodes
/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py(3257): run_cell_async
/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py(78): _pseudo_sync_runner
/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py(3030): _run_cell
/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py(2975): run_cell
/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py(539): run_cell
/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py(302): do_execute
/usr/local/lib/python3.10/dist-packages/tornado/gen.py(234): wrapper
/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py(539): execute_request
/usr/local/lib/python3.10/dist-packages/tornado/gen.py(234): wrapper
/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py(261): dispatch_shell
/usr/local/lib/python3.10/dist-packages/tornado/gen.py(234): wrapper
/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py(361): process_one
/usr/local/lib/python3.10/dist-packages/tornado/gen.py(786): run
/usr/local/lib/python3.10/dist-packages/tornado/gen.py(825): inner
/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py(738): _run_callback
/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py(685): <lambda>
/usr/lib/python3.10/asyncio/events.py(80): _run
/usr/lib/python3.10/asyncio/base_events.py(1909): _run_once
/usr/lib/python3.10/asyncio/base_events.py(603): run_forever
/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py(195): start
/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py(619): start
/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py(992): launch_instance
/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py(37): <module>
/usr/lib/python3.10/runpy.py(86): _run_code
/usr/lib/python3.10/runpy.py(196): _run_module_as_main