Search code examples
pythonpytorchtensorscalar

How to multiply a tensor row-wise by a vector in PyTorch?


When I have a tensor m of shape [12, 10] and a vector s of scalars with shape [12], how can I multiply each row of m with the corresponding scalar in s?


Solution

  • You need to add a corresponding singleton dimension:

    m * s[:, None]
    

    s[:, None] has size of (12, 1) when multiplying a (12, 10) tensor by a (12, 1) tensor pytorch knows to broadcast s along the second singleton dimension and perform the "element-wise" product correctly.