pytorch条件选择行和列

在使用pytorch中，我们有时会需要按照某些条件去选择行和列（甚至是任意一维），这时候可以这么选择：

a=torch.tensor(……) # a是一个n维的tensor
b=a[:,:,…,condition,…:,:] # 在某维上指定条件

比方说：

>>> a=torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
>>> condition=[True,False,True]
>>> a[condition,:]  #选择满足条件的行
tensor([[1, 2, 3],
        [7, 8, 9]])
>>> a[:,condition] #选择满足条件的列
tensor([[1, 3],
        [4, 6],
        [7, 9]])