Solved by Kamer_Ali_Yuksel in post #2 
              
              
                I have managed to do it via index_copy_ , is this the right solution? 
opt_weights.data.index_copy_(0, (tr == kth).nonzero().flatten(), dist.sample_n(opt_weights[tr == kth].shape[0]))