Load a tensor by name and return it as a candle Tensor on the given device.
(
&self,
name: &str,
dtype: candle_core::DType,
device: &candle_core::Device,
)
| 268 | |
| 269 | /// Load a tensor by name and return it as a candle Tensor on the given device. |
| 270 | pub fn load_tensor( |
| 271 | &self, |
| 272 | name: &str, |
| 273 | dtype: candle_core::DType, |
| 274 | device: &candle_core::Device, |
| 275 | ) -> Result<candle_core::Tensor> { |
| 276 | let data = self.read_tensor(name)?; |
| 277 | let tensor = candle_core::Tensor::from_raw_buffer( |
| 278 | &data.bytes, |
| 279 | data.dtype, |
| 280 | &data.shape, |
| 281 | &candle_core::Device::Cpu, |
| 282 | ).map_err(|e| anyhow::anyhow!("from_raw_buffer({name}): {e}"))?; |
| 283 | let tensor = if tensor.dtype() != dtype { |
| 284 | tensor.to_dtype(dtype).map_err(|e| anyhow::anyhow!("to_dtype({name}): {e}"))? |
| 285 | } else { |
| 286 | tensor |
| 287 | }; |
| 288 | if !device.is_cpu() { |
| 289 | tensor.to_device(device).map_err(|e| anyhow::anyhow!("to_device({name}): {e}")) |
| 290 | } else { |
| 291 | Ok(tensor) |
| 292 | } |
| 293 | } |
| 294 | |
| 295 | /// Load all tensors as a HashMap suitable for VarBuilder::from_tensors(). |
| 296 | pub fn load_all( |