Struct PyTorchCheckpoint 

pub struct PyTorchCheckpoint {
    pub state_dict: StateDict,
    pub optimizer_state: Option<OptimizerState>,
    pub epoch: Option<usize>,
    pub loss_history: Option<Vec<f32>>,
    pub metadata: HashMap<String, String>,
}

PyTorch checkpoint structure.

Contains the model state_dict and optional optimizer state, epoch information, and other training metadata commonly saved in PyTorch checkpoints.

Fields

state_dict: StateDict

Model state dictionary

optimizer_state: Option<OptimizerState>

Optimizer state (if saved)

epoch: Option<usize>

Training epoch (if saved)

loss_history: Option<Vec<f32>>

Training loss history (if saved)

metadata: HashMap<String, String>

Custom metadata

Implementations

impl PyTorchCheckpoint

pub fn load<P: AsRef<Path>>(path: P) -> Result<Self>

Load a PyTorch checkpoint from a file.

Security Note

This uses pickle deserialization which can be unsafe with untrusted files. Only load checkpoints from trusted sources.

pub fn metadata(&self) -> CheckpointMetadata

Get checkpoint metadata.

pub fn state_dict(&self) -> &StateDict

Get reference to state dict.

pub fn to_safetensors(&self) -> Result<Vec<u8>>

Convert checkpoint to Safetensors format.

This provides a safe, efficient format for storing model weights.

pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<()>

Save checkpoint in PyTorch format.

Note: This creates a simplified pickle format compatible with PyTorch.

pub fn new() -> Self

Create a new empty checkpoint.

pub fn add_tensor(&mut self, name: String, tensor: TensorData)

Add a tensor to the state dict.

pub fn set_epoch(&mut self, epoch: usize)

Set the epoch.

pub fn add_metadata(&mut self, key: String, value: String)

Add metadata entry.

Trait Implementations

impl Clone for PyTorchCheckpoint

fn clone(&self) -> PyTorchCheckpoint

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for PyTorchCheckpoint

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for PyTorchCheckpoint

fn default() -> Self

Returns the “default value” for a type.

impl<'de> Deserialize<'de> for PyTorchCheckpoint

fn deserialize<__D>(__deserializer: __D) -> Result<Self, __D::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl Serialize for PyTorchCheckpoint

fn serialize<__S>(&self, __serializer: __S) -> Result<__S::Ok, __S::Error>

where __S: Serializer,

Serialize this value into the given Serde serializer.