Allocate a chunk of memory from the buffer to tensor and copy the values.
(self, tensor)
| 88 | return self._start |
| 89 | |
| 90 | def add(self, tensor): |
| 91 | """Allocate a chunk of memory from the buffer to tensor and copy |
| 92 | the values.""" |
| 93 | assert ( |
| 94 | tensor.dtype == self.dtype |
| 95 | ), "Input tensor type {} different from buffer type {}".format( |
| 96 | tensor.dtype, self.dtype |
| 97 | ) |
| 98 | # Number of elements of the input tensor. |
| 99 | tensor_numel = torch.numel(tensor) |
| 100 | new_start = self._start + tensor_numel |
| 101 | assert ( |
| 102 | new_start <= self.numel |
| 103 | ), "Not enough memory left in the buffer ({} > {})".format( |
| 104 | tensor_numel, self.numel - self._start |
| 105 | ) |
| 106 | # New tensor is a view into the memory. |
| 107 | new_tensor = self.data[self._start : new_start] |
| 108 | self._start = new_start |
| 109 | new_tensor = new_tensor.view(tensor.shape) |
| 110 | new_tensor.copy_(tensor) |
| 111 | # Return a pointer to the new tensor. |
| 112 | return new_tensor |
| 113 | |
| 114 | def get_data(self): |
| 115 | """Return the data currently in use.""" |