torch.masked_select太好用了

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])

posted @ 2021-12-05 17:53  zae  阅读(787)  评论(0编辑  收藏  举报