Updating safetensors depe ndency

coreylowman · Sep 5, 2023 · 761c2eb · 761c2eb
1 parent 1255690
commit 761c2eb
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 12 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -4,7 +4,7 @@ resolver = "2"
 
 [workspace.dependencies]
 num-traits = { version = "0.2.15", default-features = false }
-safetensors = { version = "0.3", default-features = false }
+safetensors = { version = "0.3.3", default-features = false }
 memmap2 = { version = "0.5", default-features = false }
 rand = { version = "0.8.5", default-features = false, features = ["std_rng"] }
 rand_distr = { version = "0.4.3", default-features = false }

diff --git a/dfdx-nn-core/src/lib.rs b/dfdx-nn-core/src/lib.rs
@@ -150,17 +150,12 @@ pub trait SaveSafeTensors {
     ) -> Result<(), safetensors::SafeTensorError> {
         let mut tensors = Vec::new();
         self.write_safetensors("", &mut tensors);
-        let data = tensors
-            .iter()
-            .map(|(k, dtype, shape, data)| {
-                (
-                    k.clone(),
-                    safetensors::tensor::TensorView::new(dtype.clone(), shape.clone(), data)
-                        .unwrap(),
-                )
-            })
-            .collect::<Vec<_>>();
-        let data = data.iter().map(|i| (i.0.clone(), &i.1)).collect::<Vec<_>>();
+        let data = tensors.iter().map(|(k, dtype, shape, data)| {
+            (
+                k.clone(),
+                safetensors::tensor::TensorView::new(dtype.clone(), shape.clone(), data).unwrap(),
+            )
+        });
 
         safetensors::serialize_to_file(data, &None, path.as_ref())
     }