| 327 | |
| 328 | |
| 329 | class NestedTensor(object): |
| 330 | def __init__(self, tensors, mask: Optional[Tensor]): |
| 331 | self.tensors = tensors |
| 332 | self.mask = mask |
| 333 | |
| 334 | def to(self, device, non_blocking=False): |
| 335 | # type: (Device) -> NestedTensor # noqa |
| 336 | cast_tensor = self.tensors.to(device, non_blocking=non_blocking) |
| 337 | mask = self.mask |
| 338 | if mask is not None: |
| 339 | assert mask is not None |
| 340 | cast_mask = mask.to(device, non_blocking=non_blocking) |
| 341 | else: |
| 342 | cast_mask = None |
| 343 | return NestedTensor(cast_tensor, cast_mask) |
| 344 | |
| 345 | def record_stream(self, *args, **kwargs): |
| 346 | self.tensors.record_stream(*args, **kwargs) |
| 347 | if self.mask is not None: |
| 348 | self.mask.record_stream(*args, **kwargs) |
| 349 | |
| 350 | def decompose(self): |
| 351 | return self.tensors, self.mask |
| 352 | |
| 353 | def __repr__(self): |
| 354 | return str(self.tensors) |
| 355 | |
| 356 | |
| 357 | def setup_for_distributed(is_master): |
no outgoing calls