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