How to check the shape of a PyTorch tensor?

Published on Aug. 22, 2023, 12:18 p.m.

To check the shape of a PyTorch tensor, you can use the .shape or .size() attribute. Here is an example:

import torch

# Create a PyTorch tensor
x = torch.tensor([[1, 2], [3, 4]])

# Get the shape of the tensor using .shape
print(x.shape)

# Get the shape of the tensor using .size()
print(x.size())

In this code, we create a PyTorch tensor x and then use .shape and .size() to get its shape. Both of these attributes return a tuple of integers representing the size of each dimension of the tensor.

Alternatively, you can use the .ndim attribute to get the number of dimensions of the tensor:

# Get the number of dimensions of the tensor
print(x.ndim)

This will output 2, which is the number of dimensions in the tensor.

Tags: