[chore] add callback

hpcaitech · Oct 29, 2024 · c615a13 · c615a13
1 parent 81b5c56
commit c615a13
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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