pytorch中的部分张量操作：输出满足条件的张量与合并多维张量

import torch

a = torch.rand((6,1024,9))
# print(a)
# print(a[1, :])

# 暴力法输出满足条件的张量，但时空复杂度均高
for i in range(0,6):
     aa = a[i, :, 0]
     for j in range(len(aa)):
         if aa[j]>0.5:
             print(a[i,j,:])
         else: continue

# 在最后一维上合并一个三维tensor与一个二维tensor
b = torch.ones((6,1024))
print(b)
b1 = torch.unsqueeze(b,dim=2)
print(b1)
c = torch.cat((a,b1), dim=2)
print(c)