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

nbansal90: <p>Hello all,</p>
<p>I am getting an error during the execution of <code>torch.jit.trace</code> function for my model.  The error which I am seeing is :</p>
<pre><code class="lang-auto">[ 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
</code></pre>
<p>What could be the probable issue here ?</p>
<p>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.</p>
<p>Thanks,<br>
Nitin Bansal</p>

ptrblck: <p>Could you post a minimal and executable code snippet to reproduce this issue, as I see another <code>RuntimeError</code> if I try to return a plain <code>dict</code>:</p>
<pre><code class="lang-python">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.
</code></pre>
<p>without using <code>strict=False</code>.</p>

nbansal90: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="171523">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<pre><code class="lang-auto">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)
</code></pre>
</blockquote>
</aside>
<p>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.</p>
<p>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 <a href="https://github.com/pytorch/pytorch/issues/61397" rel="noopener nofollow ugc">this</a>.</p>
<p>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?</p>
<p>Nitin</p>

nbansal90: <p>I think, I got the issue <a class="mention" href="/u/ptrblck">@ptrblck</a>! 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.</p>
<p>Nitin Bansal</p>

ptrblck: <p>Good to hear you’ve solved the issue and thanks for sharing the solution! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>