Search code examples
c++pytorchtorchlibtorchtorchscript

Get value from c10::Dict<c10::IValue, c10::IValue> in Pytorch C++

I'm using a TorchScript Model on Pytorch C++ Frontend.

The model in Python returns an output dict as Dict[str, List[torch.Tensor]].

When I use it in C++, it returns a c10::Dict<c10::IValue, c10::IValue>. What is the equivalent of this Python code:

value_a = output['key_a']
value_b = output['key_b']

in C++ to get value from c10::Dict?

I have tried this but it does not work.

torch::IValue key_a("key_a");
torch::IValue key_b("key_b");
c10::IValue value_a = output[key_a];
c10::IValue value_b = output[key_b];
std::cout << value_a << std::endl;
std::cout << value_b << std::endl;

And error:

error: type 'c10::Dict<c10::IValue, c10::IValue>' does not provide a subscript operator
Solution

  • You can find header file of c10:Dict here. What you want is at method (defined here), so:

    auto value_a = output.at(key_a);

    Should do the trick.

    Also you don't have to create torch::IValue key_ay("key_a") explicitly, this should be sufficient:

    auto value_a = output.at("key_a");