CODEWITHSUNDEEP
c++pytorchlibtorch
user15139744
user15139744

Reputation:

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.

Upvotes: 6

Views: 16299

Answers (3)

&#205;ndio
&#205;ndio

Reputation: 611

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

Upvotes: 9

tbhaxor
tbhaxor

Reputation: 1943

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

Upvotes: 0

Proko
Proko

Reputation: 2011

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

Upvotes: 10

Related Questions