torch.sum(),dim=0,dim=1

dim=0:(行归一)纵向压缩

dim=1:  (列归一)横向压缩

tensor([[ 1,  2,  3],
        [ 4,  5,  6]])
# 2行3列,没毛病

#dim=0 纵向压缩,保留行
b = torch.sum(a,dim=0)
print(b)

输出:tensor([ 5,  7,  9])


#dim=1,横向压缩,保留列
c = torch.sum(a,dim=1)
print(c)

输出:tensor([  6,  15])



猜你喜欢

转载自blog.csdn.net/jialibang/article/details/107985266
今日推荐