Torch: Update tensor with non-zero elements

Question

suppose I have:

>>> a = torch.tensor([1, 2, 3, 0, 0, 1])
>>> b = torch.tensor([0, 1, 3, 3, 0, 0])

I want to update b with elements in a if it's not zero. How can I beneficently do that?

Expected:

>>> b = torch.tensor([1, 2, 3, 3, 0, 1])

Answer 1

To add to the previous answer and for more simplicity you can do it by one line of code:

b = torch.where(a!=0,a, b)

Output:

tensor([1, 2, 3, 3, 0, 1])