CODEWITHSUNDEEP
c++pytorchlibtorch
Ellenys
Ellenys

Reputation: 31

PyTorch C++ FrontEnd returning multiple Tensors in forward

I was wonder how can I return a std::vector<torch::Tensor> in my forward pass of a Module Class, I read about the Macro of FORWARD_HAS_DEFAULT_ARGS in the docs, but didn’t really understand how to use it, and also how to use it for making it possible to return a vector in return. Thank you in advance.

Upvotes: 2

Views: 3630

Answers (1)

Szymon Maszke
Szymon Maszke

Reputation: 24681

FORWARD_HAS_DEFAULT_ARGS is a C++ macro and according to documentation:

This macro enables a module with default arguments in its forward method to be used in a Sequential module.

So it's not what you are after.

I assume you are returning multiple torch::Tensor values contained in std::vector. You could just do that, but you should appropriately unpack it after returning like this:

# Interprets returned IValue as your desired return type
# You may have to use module.forward(inputs) depending how you loaded model
auto outputs = module->forward(inputs).toTensorVector();
# Print first tensor
std::cout << outputs[0] << std::endl;

If you want to return multiple values of different types from forward method you should just return std::tuple containing your desired types.

After this you can unpack it like this (for two torch::Tensor return values) (source here):

auto outputs = module->forward(inputs).toTuple();
torch::Tensor out1 = outputs->elements()[0].toTensor();
torch::Tensor out2 = outputs->elements()[1].toTensor();

You could also concatenate pytorch tensors (if that's all you are returning and they are of the same shape) and use view or a-like methods to unpack it. C++ frontend is pretty similar to Python's all in all, refer to docs if in doubt.

Upvotes: 3

Related Questions