How to get torch::Tensor shape

Question

If we << a torch::Tensor

#include <torch/script.h>
int main()
{
    torch::Tensor input_torch = torch::zeros({2, 3, 4});
    std::cout << input_torch << std::endl;
    return 0;
}

we see

(1,.,.) = 
  0  0  0  0
  0  0  0  0
  0  0  0  0

(2,.,.) = 
  0  0  0  0
  0  0  0  0
  0  0  0  0
[ CPUFloatType{2,3,4} ]

How to get the tensor shape (that 2,3,4)? I searched https://pytorch.org/cppdocs/api/classat_1_1_tensor.html?highlight=tensor for an API call but couldn't find one. And I searched for the operator<< overload code, and also couldn't find it.

Photochemistry

Answer 1

You can use torch::sizes() method

IntArrayRef sizes()

It's equivalent of shape in python. Furthermore you can access specific size at given ax (dimension) by invoking torch::size(dim). Both functions are in the API page you linked

Rotarian

shall be used as: my_tensor.sizes() [commenting this because I tried with torch::sizes(my_tensor) and ... it failed] – Marti

Answer 2

What works for me is:

#include <torch/script.h>
int main()
{
    torch::Tensor input_torch = torch::zeros({2, 3, 4});

    std::cout << "dim 0: " << input_torch.sizes()[0] << std::endl;
    std::cout << "dim 1: " << input_torch.sizes()[1] << std::endl;
    std::cout << "dim 2: " << input_torch.sizes()[2] << std::endl;

    assert(input_torch.sizes()[0]==2);
    assert(input_torch.sizes()[1]==3);
    assert(input_torch.sizes()[2]==4);
    
    return 0;
}

Platform:

libtorch 1.11.0
CUDA 11.3

Corpse

Answer 3

Well i have been using torch::_shape_as_tensor(tensor) which gives you another tensor object

Heathenish