[tmp] add write_tensor

2024-10-23 06:25:53 +00:00 · 2024-10-23 06:25:53 +00:00 · ff14144d9c
parent ad6558e91c
commit ff14144d9c
1 changed files with 1 additions and 1 deletions
--- 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)