diff --git a/tensornvme/async_file_io.py b/tensornvme/async_file_io.py index cdf52dd..5ee073b 100644 --- a/tensornvme/async_file_io.py +++ b/tensornvme/async_file_io.py @@ -40,8 +40,9 @@ def save( self.write(header_bytes) for tensor in tensors: - tensor_bytes = _tobytes(tensor, '') - self.write(tensor_bytes) + self.io.write(tensor.data_ptr(), tensor.numel() * tensor.element_size(), self.offset) + self.offset += tensor.numel() * tensor.element_size() + self.buffers.append(tensor) def flush(self) -> None: pass