Hi xdwang!



 xdwang0726:

Example:
a = torch.tensor([[[1,2,3],[5,6,7]], [[1,3,5],[5,8,7]], [[1,1,5],[5,8,3]]]), size(2, 2, 3)
b = torch.tensor([[1,0,1], [1, 1, 0], [0, 1, 0]]), size(2, 3)
(where the first dimension is batch size)
expected output:
torch.tensor([[[1, 0, 3], [5, 0, 7]], [[1, 3, 0], [5, 8, 0]], [[0, 1, 0], [0, 8, 0]]) , size(2, 2, 3)


Yes.
You can use broadcasting 1 together with element-wise multiplication, *,
provided you use unsqueeze() so that the dimensions line up correctly:
>>> import torch
>>> torch.__version__
'1.9.0'
>>> a = torch.tensor([[[1,2,3],[5,6,7]], [[1,3,5],[5,8,7]], [[1,1,5],[5,8,3]]])
>>> b = torch.tensor([[1,0,1], [1, 1, 0], [0, 1, 0]])
>>> a * b.unsqueeze (1)
tensor([[[1, 0, 3],
         [5, 0, 7]],

        [[1, 3, 0],
         [5, 8, 0]],

        [[0, 1, 0],
         [0, 8, 0]]])

(As an aside, the sizes you give are wrong; your batch size is 3 (not 2).)
Best.
K. Frank