It was used multiple time but when you evaluate the tensor that represents the loss (KLqp) the computation is sync (not completely sure about this, so please verify it on your side).
You are right, I was thinking about an indexing mask but not a 0-1 mask. Yes in that case you do need to also index data first.