diff options
| author | Nicolas Patry <patry.nicolas@protonmail.com> | 2023-08-09 16:50:11 +0200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-08-09 16:50:11 +0200 |
| commit | dece0b8a76c5e816cf93013f2ee54fd6e2bcbcae (patch) | |
| tree | db22146ca8d4bacd9cdc76672f87d949d84cb36e /candle-core/src | |
| parent | b80348d22f8f0dadb6cc4101bde031d5de69a9a5 (diff) | |
| parent | dba31473d40c88fed22574ba96021dc59f25f3f7 (diff) | |
| download | candle-dece0b8a76c5e816cf93013f2ee54fd6e2bcbcae.tar.gz candle-dece0b8a76c5e816cf93013f2ee54fd6e2bcbcae.tar.bz2 candle-dece0b8a76c5e816cf93013f2ee54fd6e2bcbcae.zip | |
Merge pull request #263 from huggingface/book_3
Book 3 (advanced loading + hub)
Diffstat (limited to 'candle-core/src')
| -rw-r--r-- | candle-core/src/safetensors.rs | 6 |
1 files changed, 5 insertions, 1 deletions
diff --git a/candle-core/src/safetensors.rs b/candle-core/src/safetensors.rs index 1880a041..132fb914 100644 --- a/candle-core/src/safetensors.rs +++ b/candle-core/src/safetensors.rs @@ -242,7 +242,11 @@ fn convert_back(tensor: &Tensor) -> Result<Vec<u8>> { pub fn load<P: AsRef<Path>>(filename: P, device: &Device) -> Result<HashMap<String, Tensor>> { let data = std::fs::read(filename.as_ref())?; - let st = safetensors::SafeTensors::deserialize(&data)?; + load_buffer(&data[..], device) +} + +pub fn load_buffer(data: &[u8], device: &Device) -> Result<HashMap<String, Tensor>> { + let st = safetensors::SafeTensors::deserialize(data)?; st.tensors() .into_iter() .map(|(name, view)| Ok((name, view.load(device)?))) |