| 40 | return maxes |
| 41 | |
| 42 | class NestedTensor(object): |
| 43 | def __init__(self, tensors, mask: Optional[Tensor]): |
| 44 | self.tensors = tensors |
| 45 | self.mask = mask |
| 46 | |
| 47 | def to(self, device): |
| 48 | # type: (Device) -> NestedTensor # noqa |
| 49 | cast_tensor = self.tensors.to(device) |
| 50 | mask = self.mask |
| 51 | if mask is not None: |
| 52 | assert mask is not None |
| 53 | cast_mask = mask.to(device) |
| 54 | else: |
| 55 | cast_mask = None |
| 56 | return NestedTensor(cast_tensor, cast_mask) |
| 57 | |
| 58 | def decompose(self): |
| 59 | return self.tensors, self.mask |
| 60 | |
| 61 | def __repr__(self): |
| 62 | return str(self.tensors) |
| 63 | |
| 64 | def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): |
| 65 | # TODO make this more general |
no outgoing calls
no test coverage detected