Expand a tensor with zeros

1

Given an image tensor with a shape of:
(1,3,640,480)

I want to expand the image tensor to a shape of:
(1,3,640,640)

I want to fill the newly added space with zeroes.

Here’s an example of the desired result:

# Given tensor shape (2,2)
>>> img = torch.arange(100, 104).view(2,2)
tensor([[100, 101],
        [102, 103]])

>>> dst = torch.zeros(3,3).long()
tensor([[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]])

# Desired result...
# Make image shape (3,3)
>>> dst.put_(torch.tensor((0,1,3,4)), img)
tensor([[100, 101,   0],
        [102, 103,   0],
        [  0,   0,   0]])   # < zeros
#                    ^ zeros

# Notice the image is now expanded 
# with zeros along the right and bottom.

How do I expand an image shape (1,3,640,480) into a shape (1,3,640,640)?

Thank you :slightly_smiling_face:

Solution:

dst[..., :640, :480] = img

New Problem:

I ran the above expression through torch.jit.trace and discovered that tracing fails. Is there a torch function that can achieve the desired result (expand zeros on a tensor) while complying with torch.jit.trace?