I have a tensor A with shape of (b, 1) and a tensor B with shape of (b, h, w). The question is how could I perform this process quickly by using more efficient codes?

for i in range(b):
B[i, int(A[i]):] = 0

That is, I want to set part of B (a bunch of matrices) to zero, and the beginning row indices are in A.