How to Get the Shape of a Tensor as a List of int in Pytorch?

To get the shape of a tensor as a list in PyTorch, we can use two approaches. One using the size() method and another by using the shape attribute of a tensor in PyTorch. In this short article, we are going to see how to use both of the approaches.

Using size() method:

The size() method returns the size of the self tensor. The returned value is a subclass of a tuple.

Python3

import torch
torch.empty(3, 4, 5).size()

Output: 

torch.Size([3, 4, 5])

We cast this into a list using the list() method. 

Example:

Python3

v = torch.tensor([[1,0],[0,1]])
x = list(v.size())
print(x)

Output:

[2, 2] 

You can also use the Python interactive session as shown below:

Using shape attribute:

The tensor.shape is an alias to tensor.size(), though the shape is an attribute, and size() is a method. To verify this we can run the following in the interactive session.

We can similarly cast this into a list using the list() method. 

Example:

Python3

import torch
v = torch.tensor([[1,0],[0,1]])
x = list(v.shape)
print(x)

 Output:

 [2,2]

Alternatively using the interactive session: