[pytorch][原创]获取torch tensor中指定条件的数据

代码场景：在目标检测中获取到了类别，置信度和bbox的tensor，但是我想只想要指定类别的数据集咋办，比如只要class=1的数据，于是便有了下面代码

import torch

cls = torch.Tensor([1, 1, 2, 3, 1]) 
conf = torch.Tensor([0.1, 0.2, 0.3, 0.4, 0.5])
bbox = torch.Tensor([[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4], [5, 5, 5, 5]])
#print(bbox)
print(conf [cls==1])
print(bbox [cls==1])