Torch JIT trace problem

rayleichen: <p>The code is as follows:</p>
<pre><code class="lang-auto">class DemoTargetObjectFeatureProcessor(torch.nn.Module):
    def __init__(self):
        super(DemoTargetObjectFeatureProcessor, self).__init__()

    def forward(self, target_object_inputs):
        target_object_size = target_object_inputs[:, 0] * target_object_inputs[:, 1]
        target_object_size = target_object_size.unsqueeze(-1)
        return torch.cat((target_object_inputs, target_object_size), dim=1)

    @torch.jit.export
    def forward_1(self, target_object_inputs_dict):
        target_object_inputs = target_object_inputs_dict['a']
        return self.forward(target_object_inputs)

target_object_inputs = torch.tensor([[1,1], [2,2]])
target_object_inputs_dict = {
    "a": torch.tensor([[1,1], [2,2]])
}

fp = DemoTargetObjectFeatureProcessor()
fp.forward(target_object_inputs)
fp.forward_1(target_object_inputs_dict)

module = torch.jit.trace(fp, target_object_inputs)
module = torch.jit.trace(fp.forward, target_object_inputs)

module = torch.jit.trace(fp.forward_1, target_object_inputs_dict)
</code></pre>
<p>I’m trying to pass a named dictionary of tensors into the TorchScript module from “module = torch.jit.trace(fp.forward_1, target_object_inputs_dict)”, but I will get the error as follows:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "test_trace.py", line 50, in &lt;module&gt;
    module = torch.jit.trace(fp.forward_1, target_object_inputs_dict)
  File "/home/lei.chen/.pyenv/versions/3.6.6/lib/python3.6/site-packages/torch/jit/__init__.py", line 893, in trace
    raise AttributeError("trace doesn't support compiling individual module's functions.\n"
AttributeError: trace doesn't support compiling individual module's functions.
Please use trace_module
</code></pre>
<p>If i’m trying to use</p>
<pre><code class="lang-auto">module = torch.jit.script(fp)
module.forward_1(target_object_inputs_dict)
</code></pre>
<p>I will get the follow error:</p>
<pre><code class="lang-auto">RuntimeError:
Unsupported operation: indexing tensor with unsupported index type 'str'. Only ints, slices, and tensors are supported:
  File "test_trace.py", line 36
    @torch.jit.export
    def forward_1(self, target_object_inputs_dict):
        target_object_inputs = target_object_inputs_dict['a']
                               ~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
        return self.forward(target_object_inputs)
</code></pre>
<p>How can we pass a named dictionary of tensors into TorchScript? I want to run model inference in C++ backend using the exported TorchScript. From the PyTorch unit tests <a href="https://github.com/pytorch/pytorch/blob/5136ed0e44c65cb3747a1f22f77ccf09d54c125c/test/cpp/api/jit.cpp#L73" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/5136ed0e44c65cb3747a1f22f77ccf09d54c125c/test/cpp/api/jit.cpp#L73</a>, it seems that we can pass a c10::Dict into TorchScript module.</p>

rayleichen: <p>Any TorchScript expert can help here? I also find this issue in PyTorch github page: <a href="https://github.com/pytorch/pytorch/issues/16847" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/16847</a>. It’s claimed that “Closing as we support this in the tracer now.”, but many ppls experienced the error from dictionary as input tensors.</p>

Dipam_Vasani: <p>Hi, if you read the error, it infers the input as a tensor. You need to explicitly specify it’s a dict with annotation</p>
<pre><code class="lang-auto">from typing import Dict
import torch

def forward_1(self, target_object_inputs_dict: Dict[str, torch.Tensor]):
</code></pre>