如何检查某个数字是否在 Pytorch 张量中？

Question

对于 Pytorch 张量A：

A = tensor([1,0,0],
           [0,0,0])

有什么方法可以检查数字 1 是否是张量的一个元素A？就像是否有一个pytorch函数返回True是1是一个元素A，如果1不是一个元素，则返回False A？

谢谢，

Answer 1

torch.Tensor实施__contains__. 所以，你可以使用：

1 in A

True如果元素1在A，False则返回，否则返回。