diffusion_rs_common::core::pickle

Function read_all_with_key

source 
pub fn read_all_with_key<P: AsRef<Path>>(
    path: P,
    key: Option<&str>,
) -> Result<Vec<(String, Tensor)>>
Expand description

Read all the tensors from a PyTorch pth file with a given key.

ยงArguments

  • path - Path to the pth file.
  • key - Optional key to retrieve state_dict from the pth file. Sometimes the pth file contains multiple objects and the state_dict is the one we are interested in.