Speedup ShardedSafeTensors to load Tensors with default hints

huggingface · Nov 28, 2023 · cc816b8 · cc816b8
1 parent 7c3cfd1
commit cc816b8
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/candle-nn/src/var_builder.rs b/candle-nn/src/var_builder.rs
@@ -535,12 +535,17 @@ impl Backend for ShardedSafeTensors {
 
     fn get(
         &self,
-        _target_shape: Shape, // The size is not checked for ShardedTensors
+        target_shape: Shape, // The size is not checked for ShardedTensors
         path: &str,
         h: Self::Hints,
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
+        if h == Default::default() {
+            // no sharding
+            return SimpleBackend::get(&self.0, target_shape, path, Default::default(), dtype, dev);
+        }
+
         let Shard {
             dim,
             rank,