Batched elementwise multiplication question

I have input_tensor of shape [batch, num_samples, num_features]

And attention_tensor of shape [batch, num_samples]

Now I need, for every batch, for every sample, multiply elementwise value from attention_tensor to each feature value in input_tensor so output tensor will have shape [batch, num_samples, num_features]

How can I do this in Pytorch with minimum overhead?

Does input_tensor * attention_tensor.unsqueeze(-1) work?

1 Like

Yes it works!
As well as input_tensor * attention_tensor.unsqueeze(2) - what I figured out by experiment.

Although I am not sure if this solution is optimal.