huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: can directly save the loaded weights (#421)

Browse Source
This commit is contained in:
Yumin Wu
2023-08-12 23:33:29 +08:00
committed by GitHub
parent 0741ebbd51
commit 16b89f5b83
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
candle-core/src/safetensors.rs
View File

@ -253,7 +253,10 @@ pub fn load_buffer(data: &[u8], device: &Device) -> Result<HashMap<String, Tenso
.collect()
}
pub fn save<P: AsRef<Path>>(tensors: &HashMap<&str, Tensor>, filename: P) -> Result<()> {
pub fn save<K: AsRef<str> + Ord + std::fmt::Display, P: AsRef<Path>>(
tensors: &HashMap<K, Tensor>,
filename: P,
) -> Result<()> {
Ok(st::serialize_to_file(tensors, &None, filename.as_ref())?)
}