Solved by yf225 in post #2 
              
              
                Does the following work? 
auto a = torch::zeros({10, 4});
auto b = torch::ones({4, 4});
auto mask = torch::tensor({0, 1, 0, 0, 1, 0, 1, 1, 0, 0}, torch::dtype(torch::kBool));
a.index_put_({mask}, b);