pub trait ApplyGradientExt<Delta, T>: ApplyGradient<Delta, T> {
type Velocity;
// Required methods
fn apply_gradient_with_momentum(
&mut self,
grad: &Delta,
lr: T,
momentum: T,
velocity: &mut Self::Velocity,
) -> Result<Self::Output>;
fn apply_gradient_with_decay_and_momentum(
&mut self,
grad: &Delta,
lr: T,
decay: T,
momentum: T,
velocity: &mut Self::Velocity,
) -> Result<Self::Output>;
}
Expand description
This trait extends the ApplyGradient trait by allowing for momentum-based optimization