find index of certain value in 3 dimension array in parallel way

Question

I made a zero_mask2 function (copied below) in the middle of my pytorch file. However, it is too slow. So, I\'m looking for a better way.

First, Let me e