Debugging JIT tracing/scripting

Hi, I’m currently trying to JIT compile CenterTrack with either tracing or scripting.

CenterTrack uses DCNv2. I’ve found some resources that mentioned I needed to fix a couple things first:

Now I’ve removed some of the initial errors I got when trying to trace, but am stuck the current ones.
When I try to trace:

# Try trace
input_var = torch.zeros([1, 3, 736, 1280], dtype=torch.float32).cuda()
try:	
    traced_script_module = torch.jit.trace(detector.model, input_var)
except Exception as e:
    print(f'Exception (trace): {e}')

The only thing that is returned is:

Exception (trace): Only tensors, lists and tuples of tensors can be output from traced functions

And when I instead try torch script I have two issues, a small one being requirements of default values for tensors, but a larger one being heavy use of *args as inputs throughout the codebase.

I suspect that the torch trace issue comes from some forward() methods returning dictionaries. Is that not supported in v1.5? Is there a recommended alternative? And also is there a way to find out where in my python code the error “Only tensors, lists and tuples of tensors can be output from traced functions” refers to?

Any help is greatly appreciated!