From ff14144d9cebc2d334de714f52d1b599ece0a116 Mon Sep 17 00:00:00 2001 From: botbw 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)