Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to get torch::Tensor shape

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.


1 Answers

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

like image 55
Proko Avatar answered Aug 28 '26 11:08

Proko