PyTorch torch.squeeze Function
PyTorch torch Reference Manual
torch.squeezeIt is a function in PyTorch used to remove dimensions of size 1 from a tensor.
Function Definition
torch.squeeze(input, dim)
Usage Example
Example
import torch
# Create a tensor with dimensions of size 1
x = torch.randn(1, 2, 1, 3, 1)
print("Original shape:", x.shape)
# Remove all dimensions of size 1
y = torch.squeeze(x)
print("Squeezed shape:", y.shape)
# Specify the dimension
z = torch.squeeze(x, dim=0)
print("After dim=0 squeeze:", z.shape)
# Create a tensor with dimensions of size 1
x = torch.randn(1, 2, 1, 3, 1)
print("Original shape:", x.shape)
# Remove all dimensions of size 1
y = torch.squeeze(x)
print("Squeezed shape:", y.shape)
# Specify the dimension
z = torch.squeeze(x, dim=0)
print("After dim=0 squeeze:", z.shape)
The output is:
原始形状: torch.Size([1, 2, 1, 3, 1]) 挤压后形状: torch.Size([2, 3]) dim=0 挤压后: torch.Size([2, 1, 3, 1])
Other Extensions