Struct TrainingLoop

pub struct TrainingLoop<A: Float, O: UnifiedOptimizer<A>> { /* private fields */ }

Training loop helper with unified API

Implementations

impl<A: Float + ScalarOperand + Debug, O: UnifiedOptimizer<A>> TrainingLoop<A, O>

pub fn new(optimizer: O) -> Self

Create a new training loop

pub fn with_scheduler( self, scheduler: Box<dyn LearningRateScheduler<A>>, ) -> Self

Add a learning rate scheduler

pub fn step<D: Dimension>( &mut self, params: &mut [Parameter<A, D>], ) -> Result<()>

Perform one training step

pub fn zero_grad<D: Dimension>(&self, params: &mut [Parameter<A, D>])

Zero gradients

pub fn get_lr(&self) -> A

Get current learning rate