From ff14144d9cebc2d334de714f52d1b599ece0a116 Mon Sep 17 00:00:00 2001
From: botbw <wang1570@e.ntu.edu.sg>
Date: Wed, 23 Oct 2024 06:25:53 +0000
Subject: [PATCH] [tmp] add write_tensor
---
colossalai/utils/safetensors.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/colossalai/utils/safetensors.py b/colossalai/utils/safetensors.py
index 9aa3558d9..390e12511 100644
--- a/colossalai/utils/safetensors.py
+++ b/colossalai/utils/safetensors.py
@@ -61,4 +61,4 @@ def save(f_writer: AsyncFileWriter, state_dict: Dict[str, torch.Tensor]) -> None
f_writer.write(header_bytes)
for tensor in tensors:
- f_writer.write_raw(tensor, tensor.data_ptr(), tensor.numel() * tensor.element_size(), f_writer.offset)
+ f_writer.write_tensor(tensor)