| 44 | self.buffer.storage().resize_(0) |
| 45 | |
| 46 | def append(self, tensor: Tensor, callback_fn: Callable): |
| 47 | tensor_size = tensor.numel() |
| 48 | offset = self.offset |
| 49 | self.buffer[offset : offset + tensor_size].copy_(tensor.flatten()) |
| 50 | self.offset += tensor_size |
| 51 | |
| 52 | # callback will be given the reduced result |
| 53 | if callback_fn is not None: |
| 54 | result_view = self.buffer[offset : offset + tensor_size].view(tensor.shape) |
| 55 | self.callbacks.append(functools.partial(callback_fn, result_view)) |
| 56 | |
| 57 | @property |
| 58 | def avail_size(self) -> int: |