Struct burn_core::optim::RmsProp

pub struct RmsProp<B: Backend> { /* private fields */ }

Optimizer that implements stochastic gradient descent with momentum. The optimizer can be configured with RmsPropConfig.

Trait Implementations

impl<B: Backend> SimpleOptimizer<B> for RmsProp<B>

type State<const D: usize> = RmsPropState<B, D>

The state of the optimizer. It also implements record, so that it can be saved.

fn step<const D: usize>( &self, lr: LearningRate, tensor: Tensor<B, D>, grad: Tensor<B, D>, state: Option<Self::State<D>> ) -> (Tensor<B, D>, Option<Self::State<D>>)

The optimizer step is performed for one tensor at a time with its gradient and state.

fn to_device<const D: usize>( state: Self::State<D>, device: &<B as Backend>::Device ) -> Self::State<D>

Change the device of the state.