How to squeeze and unsqueeze a tensor in PyTorch?

Published January 03, 2022

If you want to squeeze a tensor, you should use the torch.squeeze() method. This method returns the new tensor dimension of the input tensor. It does, however, eliminate size 1.

If you want to unsqueeze a tensor, you should use the touch.unsqueeze() function. This function returns the new tensor dimension of the size 1 placed at the specified point.

 

How to squeeze and unsqueeze a tensor in PyTorch

To squeeze and unsqueeze a tensor in PyTorch, follow the steps below:

Step 1: Importing the torch library is the initial step.

Step 2: Construct and print a tensor.

Step 3: Ascertain torch.squeeze (input). It squeezes (removes) the size 1 and returns a tensor with all of the remaining dimensions of the input tensor.

Step 4: Select torch.unsqueeze (input, dim). After adding a new dimension of size 1 at the supplied dim, it returns the tensor.

Step 5: Squeezed and/or unsqueezed tensors should be printed.

 

Example

import torch T = torch.ones(2,1,2) # size 2x1x2 print("Original Tensor T:\n", T ) print("Size of T:", T.size()) squeezed_T = torch.squeeze(T) print("Squeezed_T\n:", squeezed_T ) print("Size of Squeezed_T:", squeezed_T.size())

 

Output:

Original Tensor T: tensor([[[1., 1.]],          [[1., 1.]]]) Size of T: torch.Size([2, 1, 2]) Squeezed_T : tensor([[1., 1.],          [1., 1.]]) Size of Squeezed_T: torch.Size([2, 2])

 

 

Download Source code