File size: 791 Bytes
da50cbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | from typing import List
from PIL import Image
from torch.nn import functional as F
def pad_to_length(tensor, len, dim=1):
return F.pad(tensor, [0] * ((tensor.dim() - 2) * 2) + [0, len - tensor.shape[dim]])
def expand_as(tensor, other):
"""
Expands a tensor to match the dimensions of another tensor.
If tensor has shape [b] and other has shape [b, c, h, w],
this function will reshape tensor to [b, 1, 1, 1] to enable broadcasting.
Args:
tensor (`torch.FloatTensor`): The tensor to expand
other (`torch.FloatTensor`): The tensor whose shape will be matched
Returns:
`torch.FloatTensor`: The expanded tensor
"""
for _ in range(other.ndim - tensor.ndim):
tensor = tensor.unsqueeze(-1)
return tensor |