| 340 | |
| 341 | |
| 342 | class NestedTensor(object): |
| 343 | def __init__(self, tensors, mask: Optional[Tensor]): |
| 344 | self.tensors = tensors |
| 345 | self.mask = mask |
| 346 | |
| 347 | def to(self, device, non_blocking=False): |
| 348 | cast_tensor = self.tensors.to(device, non_blocking=non_blocking) |
| 349 | mask = self.mask |
| 350 | if mask is not None: |
| 351 | assert mask is not None |
| 352 | cast_mask = mask.to(device, non_blocking=non_blocking) |
| 353 | else: |
| 354 | cast_mask = None |
| 355 | return NestedTensor(cast_tensor, cast_mask) |
| 356 | |
| 357 | def record_stream(self, *args, **kwargs): |
| 358 | self.tensors.record_stream(*args, **kwargs) |
| 359 | if self.mask is not None: |
| 360 | self.mask.record_stream(*args, **kwargs) |
| 361 | |
| 362 | def decompose(self): |
| 363 | return self.tensors, self.mask |
| 364 | |
| 365 | def __repr__(self): |
| 366 | return str(self.tensors) |
| 367 | |
| 368 | |
| 369 | def setup_for_distributed(is_master): |
no outgoing calls
no test coverage detected