【pytorch学习】unsqueeze()函数

unsqueeze()函数

作用：增加tensor的维度；

import torch
import numpy as np

a =  torch.tensor(np.arange(0, 6))
print(a)
print(a.shape)

a = a.unsqueeze(1)
print(a)
print(a.shape)