| |
| |
| |
| |
| |
| |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| def pad_x_to_y(x, y, axis: int = -1): |
| if axis != -1: |
| raise NotImplementedError |
| inp_len = y.shape[axis] |
| output_len = x.shape[axis] |
| return nn.functional.pad(x, [0, inp_len - output_len]) |
|
|
|
|
| def shape_reconstructed(reconstructed, size): |
| if len(size) == 1: |
| return reconstructed.squeeze(0) |
| return reconstructed |
|
|
|
|
| def tensors_to_device(tensors, device): |
| """Transfer tensor, dict or list of tensors to device. |
| |
| Args: |
| tensors (:class:`torch.Tensor`): May be a single, a list or a |
| dictionary of tensors. |
| device (:class: `torch.device`): the device where to place the tensors. |
| |
| Returns: |
| Union [:class:`torch.Tensor`, list, tuple, dict]: |
| Same as input but transferred to device. |
| Goes through lists and dicts and transfers the torch.Tensor to |
| device. Leaves the rest untouched. |
| """ |
| if isinstance(tensors, torch.Tensor): |
| return tensors.to(device) |
| elif isinstance(tensors, (list, tuple)): |
| return [tensors_to_device(tens, device) for tens in tensors] |
| elif isinstance(tensors, dict): |
| for key in tensors.keys(): |
| tensors[key] = tensors_to_device(tensors[key], device) |
| return tensors |
| else: |
| return tensors |
|
|
