If you don’t want the changes of W_mat to be reflected on W_mat_all, then you can just .clone() just after the select. The gradients will flow back properly through the clone operation.
If you do want these changes to be reflected, you can work with W_mat_all[idx, cnt, indS] = 1 directly.