Torch.jit.trace fails with error, as output does not match dict format

Hello all,

I am getting an error during the execution of torch.jit.trace function for my model. The error which I am seeing is :

[ CPUFloatType{3000,4} ], label_volume_coord: [ CPULongType{0,3} ], label_volume_feats: 
[CPULongType{0} ]}) of traced region cannot be understood by the tracer, only outputs 
matchingdict[Union[str, Tensor], Union[Tensor, Tuple[Tensor, ...]]] can be a dictionary 
output of a traced function

What could be the probable issue here ?

PS: I have made sure that format of each element in the output dict is of type Tensor. which is required for jit.trace to function correctly.

Thanks,
Nitin Bansal

Could you post a minimal and executable code snippet to reproduce this issue, as I see another RuntimeError if I try to return a plain dict:

def fun(x):
    out = {"x": x}
    return out

x = torch.randn(10, 10)
out = fun(x)

fun_traced = torch.jit.trace(fun, x)
out = fun_traced(x)
# RuntimeError: Encountering a dict at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for `list`, use a `tuple` instead. for `dict`, use a `NamedTuple` instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.

without using strict=False.

I am not quite able to reproduce it using a simpler example. Although, In my case I am returning a dictionary of tensors. But when I write a simple function returning multiple 2D tensors in a dictionary, it passes the trace step.

Having said that, I did stumble upon a related issue, where i see a similar error being prompted, although I am not sure why this PR was closed, without actually spelling out what is the fix . The PR I am talking here is this.

The error message I receive is almost same as it is pointed in this problem report. I have tried reaching out to people who were working on this, but have got no response. Does this help us in zeroing down, what could be the probable issue for my case?

Nitin

I think, I got the issue @ptrblck! I am actually returning two dictionaries not one in my case. I think just need to collapse everything in to one dictionary and it should work fine. I just confirmed it using a simple function, returning two dict. Gives similar error to what I am achieving.

Nitin Bansal

Good to hear you’ve solved the issue and thanks for sharing the solution! :slight_smile: