From 5522bbc57c2967f3c8fb8fa9ab8a82d2c9ff8db8 Mon Sep 17 00:00:00 2001 From: yinqiwen Date: Mon, 1 Apr 2024 18:10:08 +0800 Subject: [PATCH] Add fn 'get_with_hints_dtype' in VarBuilder (#1877) (#1897) * quantized models(awq/squeezellm/...) have multiple data type tensors, use 'get_with_hints_dtype' to load tensors with given dtype --- candle-nn/src/var_builder.rs | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/candle-nn/src/var_builder.rs b/candle-nn/src/var_builder.rs index bf0902196e..7de4604456 100644 --- a/candle-nn/src/var_builder.rs +++ b/candle-nn/src/var_builder.rs @@ -178,16 +178,27 @@ impl<'a, B: Backend> VarBuilderArgs<'a, B> { name: &str, hints: B::Hints, ) -> Result { - let path = self.path(name); - self.data - .backend - .get(s.into(), &path, hints, self.data.dtype, &self.data.device) + self.get_with_hints_dtype(s, name, hints, self.data.dtype) } /// Retrieve the tensor associated with the given name at the current path. pub fn get>(&self, s: S, name: &str) -> Result { self.get_with_hints(s, name, Default::default()) } + + /// Retrieve the tensor associated with the given name & dtype at the current path. + pub fn get_with_hints_dtype>( + &self, + s: S, + name: &str, + hints: B::Hints, + dtype: DType, + ) -> Result { + let path = self.path(name); + self.data + .backend + .get(s.into(), &path, hints, dtype, &self.data.device) + } } struct Zeros;