| 750 | |
| 751 | |
| 752 | class NestedTensor(object): |
| 753 | def __init__(self, tensors, mask: Optional[Tensor]): |
| 754 | self.tensors = tensors |
| 755 | self.mask = mask |
| 756 | |
| 757 | def to(self, device): |
| 758 | # type: (Device) -> NestedTensor # noqa |
| 759 | cast_tensor = self.tensors.to(device) |
| 760 | mask = self.mask |
| 761 | if mask is not None: |
| 762 | assert mask is not None |
| 763 | cast_mask = mask.to(device) |
| 764 | else: |
| 765 | cast_mask = None |
| 766 | return NestedTensor(cast_tensor, cast_mask) |
| 767 | |
| 768 | def decompose(self): |
| 769 | return self.tensors, self.mask |
| 770 | |
| 771 | def cuda(self): |
| 772 | return self.to('cuda') |
| 773 | |
| 774 | def __repr__(self): |
| 775 | return str(self.tensors) |
| 776 | |
| 777 | |
| 778 | def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): |
no outgoing calls
no test coverage detected