PyTorch Tensor 去除维度为 1 的维度：torch.squeeze() 函数使用指南

可以使用 torch.squeeze() 函数来实现将维度为 1 的维度去掉。

代码如下：

import torch

# 创建一个(1, 16, 256, 256, 256)的tensor
tensor = torch.randn(1, 16, 256, 256, 256)

# 使用squeeze()函数去掉维度为1的维度
tensor_squeezed = torch.squeeze(tensor)

# 打印去掉维度为1后的tensor的形状
print(tensor_squeezed.shape)

输出结果为：

torch.Size([16, 256, 256, 256])

可以看到，维度为 1 的维度已经被去掉，tensor 的形状变为 (16, 256, 256, 256)。

torch.squeeze() 函数可以自动识别并移除 Tensor 中所有维度为 1 的维度，使代码更加简洁高效。