When you do indices[:, i] == 0 you get a mask which is a ByteTensor.
If you want to use it to perform operations with FloatTensor when you do a0 * grad_output, then you need to cast it to float: a0.float() * grad_output.