From c615a136140cf37f6e294948802147fe42d22ab1 Mon Sep 17 00:00:00 2001 From: botbw Date: Tue, 29 Oct 2024 03:35:51 +0000 Subject: [PATCH] [chore] add callback --- tensornvme/async_file_io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensornvme/async_file_io.py b/tensornvme/async_file_io.py index 625a61e..b0604f2 100644 --- a/tensornvme/async_file_io.py +++ b/tensornvme/async_file_io.py @@ -33,7 +33,7 @@ def write_raw(self, py_ref: object, buffer: int, n_bytes: int, offset: int) -> N def write_tensor(self, tensor: Tensor) -> None: self.buffers.append(tensor) - self.io.write_tensor(tensor, self.offset) # , partial(AsyncFileWriter.gc_callback, self.buffers, len(self.buffers) - 1)) + self.io.write_tensor(tensor, self.offset, partial(AsyncFileWriter.gc_callback, self.buffers, len(self.buffers) - 1)) self.offset += tensor.numel() * tensor.element_size() @staticmethod