How to mask a tensor in c++?

Let’s say I have the tensor:

[22.2, 3.3, 3.2, 10.2, 3.2]

let’s also say that i have another tensor:

[0, 3, 1, 1, 3]

I’d like to turn the first tensor to:

[22.2, 0.0, 3.2, 10.2, 0.0]

That is where the second tensor’s value is 3, i want the first tensor’s value to be reset to 0.

If someone can write the 1-2 lines of c++ that corresponds to this it’d help me a lot.

Thank for the help.

I found the answer. It’s something like this:

torch::where(t2 != 3, t1, torch::zeros_like(t1));