Apply a if operation pointwise on a torch tensor

So there is something that’s I can’t figure out, is it possible to do the operation (A==0 + A==1)>0 on the same operation when A is a pytorch tensor ?
So the operation result is simply a matrix equal to true were simple A[x] in [0,1] that does the or operation pointwise.
Same for the operation (A!=0 * A!=1) to do the and operation
Thanks in advance !

I don’t know if I understand your use case correctly, but would this work?

a = torch.randint(0, 3, (10, 10))
res = (a==0) | (a==1)