Struct burn_core::optim::momentum::Momentum

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

Momemtum implementation that transforms gradients.

Implementations

impl<B: Backend> Momentum<B>

pub fn new(config: &MomentumConfig) -> Self

Creates a new momentum from a config.

pub fn transform<const D: usize>( &self, grad: Tensor<B, D>, state: Option<MomentumState<B, D>> ) -> (Tensor<B, D>, MomentumState<B, D>)

Transforms a gradient.

Arguments

• grad - Gradient to transform.
• state - State of the optimizer.

Returns

• grad - Transformed gradient.
• state - State of the optimizer.