huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enhance pickle to retrieve state_dict with a given key (#1671)

Browse Source
This commit is contained in:
Dilshod Tadjibaev
2024-02-06 14:17:33 -06:00
committed by GitHub
parent a90fc5ca5a
commit b75e8945bc
6 changed files with 61 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
candle-nn/src/var_builder.rs
View File

@ -484,7 +484,7 @@ impl<'a> VarBuilder<'a> {
/// Initializes a `VarBuilder` that retrieves tensors stored in a pytorch pth file.
pub fn from_pth<P: AsRef<std::path::Path>>(p: P, dtype: DType, dev: &Device) -> Result<Self> {
let pth = candle::pickle::PthTensors::new(p)?;
let pth = candle::pickle::PthTensors::new(p, None)?;
Ok(Self::from_backend(Box::new(pth), dtype, dev.clone()))
}
}