| | |
| | |
| | |
| | |
| | |
| | |
| |
|
| | 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 |
| |
|
