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