Skip to content

Instantly share code, notes, and snippets.

@ultranity

ultranity/pad2shape.py

Created November 16, 2022 14:56
Show Gist options
  • Save ultranity/21b9731f272d08c5cadba67e4e4401e2 to your computer and use it in GitHub Desktop.
Save ultranity/21b9731f272d08c5cadba67e4e4401e2 to your computer and use it in GitHub Desktop.
Download ZIP
pytorch pad tensor to target shape
Raw
pad2shape.py
def pad(x:torch.Tensor, shape:torch.Size, mode='right with zero'):
p = [0]*len(shape)*2
for i, s in enumerate(shape):
p[2*(len(shape)-1-i)+1] = s - x.size(i)
return torch.nn.functional.pad(x, p)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment