torch.masked_select用于截取valid的标签太好用了:

pred=torch.tensor([2,3,4,5])
mask=torch.ByteTensor([1,1,0,0]) 
torch.masked_select(pred,mask) # deprecated, 用bool值更好 mask=torch.tensor([1,1,0,0],dtype=torch.bool)

output:
tensor([2, 3])