How to fill values of tensor by mask in C++ API?

Hello! I’m trying to google an answer for a few hours but can’t find it. So asking

How to implement this code in C++:

a = torch.zeros((10, 4))
b = torch.ones((4, 4))
mask = torch.tensor([0, 1, 0, 0, 1, 0, 1, 1, 0, 0], dtype='bool')
a[mask] = b

I am trying like this but it returns error:

  at::Tensor a = at::zeros({10, 4});
  at::Tensor b = at::ones({4, 4});
  bool _mask[]{false, true, false, false, true, false, true, true, false, false};
  at::Tensor mask = at::from_blob(_mask, {10}, torch::kBool);

first:

a.masked_fill(mask, b);

second:

a.index_fill(0, mask.nonzero(), b);

Thanks for reply!

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);
2 Likes

Yes, it works, thank you very much!