在使用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]])