github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-04-28 11:45:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

[tmp] add write_tensor

Browse Source
This commit is contained in:
botbw 2024-10-23 06:25:53 +00:00
parent ad6558e91c
commit ff14144d9c
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/utils/safetensors.py
View File

@ -61,4 +61,4 @@ def save(f_writer: AsyncFileWriter, state_dict: Dict[str, torch.Tensor]) -> None
f_writer.write(header_bytes) f_writer.write(header_bytes)
for tensor in tensors: 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)