Three dimension array value check

I have array “A” and “B” of size [1,3,10,10]. Variable B all value is 0.8 . I want to check all three dimension of A 10x10. if all are greater than 0.8 than i want to replace value 0 in 3 dimension position in variable B.

For, example A[0,0,0,0]=0.9,A[0,1,0,0]=0.95,A[0,2,0,0]=0.90 then B[0,0,0,0]=0,B[0,1,0,0]=0,B[0,2,0,0]=0 but if A[0,0,0,0]=0.75,A[0,1,0,0]=0.95,A[0,2,0,0]=0.90 then B should not be change.

B[B<A]=0 will do for each array in all dimension. I want to replace only when all three dimension is greater than 0.8.

Could you check if something like this would work or if I misunderstood your use case?

A = torch.randn(1, 3, 10, 10) * 2
B = torch.ones_like(A)

idx = (A > 0.8).all(1).unsqueeze(1).expand(-1, 3, -1, -1)
B[idx] = 0.
1 Like

Yes. Solution works for me. Thanks