Hi,
FYI, you can replace: W_mat_all.narrow(0,idx, 1).squeeze() by W_mat_all.select(0,idx).
Also narrow (and select) are inplace operation: the returned tensor share it’s storage with the original. This means that W_mat_all and W_mat share the same storage.
So when you try to modify W_mat inplace later by doing W_mat[cnt, indS] = 1, the storage is shared by 2 tensors already.