Function save_training_checkpoint 

pub fn save_training_checkpoint(
    path: &Path,
    model_weights: &HashMap<String, Tensor>,
    optimizer_state: &HashMap<String, Tensor>,
) -> Result<(), ModelError>

Save a full training checkpoint: model weights + optimizer state.

Both maps use string keys. Optimizer state keys are automatically prefixed to avoid collisions with model weight names.