[fio] use raw tensor ptr instead of numpy

hpcaitech · Oct 14, 2024 · 557f389 · 557f389
1 parent dbeb508
commit 557f389
Showing 1 changed file with 3 additions and 2 deletions.
diff --git 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