Pad several dimensions at the same time

I want to pad with zeros a tensor along different dimensions at the same time.

For example, given a tensor of shape [3,3,3,7], I want to obtain a tensor of shape [5,5,5,7], where the new entries contain a constant value (e.g., zero). For example, the value at the index [2,4,2,6] would contain 0.0.

Additionally, I would like for the method to be as efficient as possible and also preserve gradients (e.g., the element at the position [1,0,2,4] would have its gradient preserved since it was contained in the old tensor, before padding occurred).

F.pad should work:

x = torch.randn(3, 3, 3 ,7, requires_grad=True)
out = F.pad(x, (0, 0, 1, 1, 1, 1, 1, 1), mode="constant", value=0.0)

print(out.shape)
# torch.Size([5, 5, 5, 7])

print(out[2,4,2,6])
# tensor(0., grad_fn=<SelectBackward0>)

loss = out.mean().backward()
print(x.grad[1,0,2,4])
# tensor(0.0011)
1 Like