Search code examples
pythonpytorch

how sum two tensors with difference shapes_pytorch

I have two tensors with these shapes:

a_tensor : torch.Size([64, 37])
b_tensor : torch.Size([64, 300])

how can I sum them (a_tensor+b_tensor ) and pad a_tensor with zeros to be size of [64,300]?

For example, a_tensor is [[1 , 2 , 3],[3 , 2 , 1]] with shape of ([2 , 3]) and b_ensor is [[10 , 20 , 30 , 40] , [40 , 30 , 20 , 10]] with the shape of ([2,4]).

a_tensor + b_tensor would be like this:

[[1,2,3,0],[3,2,1,0]] + [[10,20,30,40],[40,30,20,10]] = [[11,22,33,40],[43,32,21,10]]
Solution

  • You can do so using nn.functional.pad:

    >>> a = torch.rand(64, 37)
>>> b = torch.rand(64, 300)

    Measure the padding amount:

    >>> r = b.size(-1) - a.size(-1)

    Pad and sum:

    >>> tf.pad(a, (0,0,r,0)) + b