Struct GradScaler 

pub struct GradScaler { /* private fields */ }

Gradient scaler for mixed-precision training

Handles loss scaling to prevent gradient underflow in fp16 training.

Implementations

impl GradScaler

pub fn new(initial_scale: f32) -> Self

Create a new gradient scaler

pub fn from_config(config: &MixedPrecisionConfig) -> Self

Create from config

pub fn scale(&self) -> f32

Get current scale

pub fn scale_loss(&self, loss: f32) -> f32

Scale a loss value

pub fn unscale_grad(&self, grad: f32) -> f32

Unscale a gradient value

pub fn unscale_and_check(&self, grads: &mut [f32]) -> bool

Unscale gradients in place and check for overflow

Returns true if gradients are valid (no overflow), false otherwise.

pub fn update(&mut self, grads_valid: bool)

Update the scale after a step

Call this after each optimizer step. Pass true if gradients were valid.

pub fn overflow_count(&self) -> usize

Get overflow count

pub fn successful_steps(&self) -> usize

Get successful step count

pub fn is_dynamic(&self) -> bool

Check if dynamic scaling is enabled

pub fn set_dynamic(&mut self, enabled: bool)

Enable/disable dynamic scaling

Trait Implementations

impl Debug for GradScaler

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for GradScaler

fn default() -> Self

Returns the “default value” for a type.