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
This commit is contained in:
yinqiwen 2024-04-01 18:10:08 +08:00 committed by GitHub
parent 888c09a3db
commit 5522bbc57c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 4 deletions

View File

@ -178,16 +178,27 @@ impl<'a, B: Backend> VarBuilderArgs<'a, B> {
name: &str,
hints: B::Hints,
) -> Result<Tensor> {
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<S: Into<Shape>>(&self, s: S, name: &str) -> Result<Tensor> {
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<S: Into<Shape>>(
&self,
s: S,
name: &str,
hints: B::Hints,
dtype: DType,
) -> Result<Tensor> {
let path = self.path(name);
self.data
.backend
.get(s.into(), &path, hints, dtype, &self.data.device)
}
}
struct Zeros;