From 16b89f5b836b9afd7883fa0bc7280116c54cb440 Mon Sep 17 00:00:00 2001
From: Yumin Wu <wuyumin@wuyumin.com>
Date: Sat, 12 Aug 2023 23:33:29 +0800
Subject: [PATCH] fix: can directly save the loaded weights (#421)

---
 candle-core/src/safetensors.rs | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/candle-core/src/safetensors.rs b/candle-core/src/safetensors.rs
index 914e5101..ec0dc766 100644
--- a/candle-core/src/safetensors.rs
+++ b/candle-core/src/safetensors.rs
@@ -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())?)
 }