From cc816b89f4a7e350625cdfc747454640c2407ef2 Mon Sep 17 00:00:00 2001
From: YISH <mokeyish@hotmail.com>
Date: Tue, 28 Nov 2023 15:38:46 +0800
Subject: [PATCH] Speedup ShardedSafeTensors to load Tensors with default hints

---
 candle-nn/src/var_builder.rs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/candle-nn/src/var_builder.rs b/candle-nn/src/var_builder.rs
index cbd238ddfa..0b266a41ee 100644
--- 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,