I am trying to trace a model which was trained using DDP. I get this error while tracing the model:
ops.profiler._record_function_exit(_1)
_517 = ^_DDPSink(<torch._C._distributed_c10d.Reducer object at 0x7fd53d9d1f70>, {'static_graph': False, 'num_iterations': 0})(image, boxes1, _516)
_518 = {"image": image, "boxes": boxes1, "density_pred": _516}
return _518
Traceback (most recent call last):
File "/content/Few-Shot-Object-Counting-PyTorch/tools/trace_model.py", line 160, in <module>
main()
File "/content/Few-Shot-Object-Counting-PyTorch/tools/trace_model.py", line 150, in main
traced_model.save(save_path)
File "/usr/local/lib/python3.10/dist-packages/torch/jit/_script.py", line 713, in save
return self._c.save(str(f), **kwargs)
RuntimeError:
Could not export Python function call '_DDPSink'. 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(506): apply
/usr/local/lib/python3.10/dist-packages/torch/nn/parallel/distributed.py(1207): forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1488): _slow_forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1501): _call_impl
/content/Few-Shot-Object-Counting-PyTorch/tools/trace_model.py(83): forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1488): _slow_forward
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1501): _call_impl
/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py(1056): trace_module
/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py(794): trace
/content/Few-Shot-Object-Counting-PyTorch/tools/trace_model.py(109): trace
/content/Few-Shot-Object-Counting-PyTorch/tools/trace_model.py(148): main
/content/Few-Shot-Object-Counting-PyTorch/tools/trace_model.py(160): <module>