Skip to main content

OptimizerState

trustformers_optim::optimizer

Trait OptimizerState 

Source 
pub trait OptimizerState {
    // Required methods
    fn zero_grad(&mut self) -> Result<()>;
    fn step(&mut self, parameters: &mut [Tensor]) -> Result<()>;
    fn get_lr(&self) -> f32;
    fn set_lr(&mut self, lr: f32);
    fn state_dict(&self) -> Result<HashMap<String, Tensor>>;
    fn load_state_dict(&mut self, state: HashMap<String, Tensor>) -> Result<()>;
}
Expand description

Trait for optimizer state management and parameter updates.

Required Methods§

Source

fn zero_grad(&mut self) -> Result<()>

Zero out gradients

Source

fn step(&mut self, parameters: &mut [Tensor]) -> Result<()>

Perform optimization step

Source

fn get_lr(&self) -> f32

Get current learning rate

Source

fn set_lr(&mut self, lr: f32)

Set learning rate

Source

fn state_dict(&self) -> Result<HashMap<String, Tensor>>

Save optimizer state to dictionary

Source

fn load_state_dict(&mut self, state: HashMap<String, Tensor>) -> Result<()>

Load optimizer state from dictionary

Implementors§