Struct GradScaler 

pub struct GradScaler { /* private fields */ }

Gradient scaler for mixed precision training.

Scales the loss to prevent gradient underflow when using F16, then unscales gradients before the optimizer step.

The scale is automatically adjusted based on whether gradients overflow.

Implementations

impl GradScaler

pub fn new() -> Self

Creates a new gradient scaler with default settings.

Default configuration:

• Initial scale: 65536.0 (2^16)
• Growth factor: 2.0
• Backoff factor: 0.5
• Growth interval: 2000 steps

pub fn with_scale(init_scale: f32) -> Self

Creates a gradient scaler with custom initial scale.

pub fn with_options( init_scale: f32, growth_factor: f32, backoff_factor: f32, growth_interval: usize, ) -> Self

Creates a gradient scaler with all custom settings.

pub fn growth_factor(self, factor: f32) -> Self

Builder: set growth factor

pub fn backoff_factor(self, factor: f32) -> Self

Builder: set backoff factor

pub fn growth_interval(self, interval: usize) -> Self

Builder: set growth interval

pub fn enabled(self, enabled: bool) -> Self

Builder: set enabled state

pub fn get_scale(&self) -> f32

Returns the current scale factor.

pub fn set_scale(&mut self, scale: f32)

Sets the scale factor.

pub fn is_enabled(&self) -> bool

Returns whether the scaler is enabled.

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

Enables or disables the scaler.

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

Scales a loss value for backward pass.

Multiply the loss by this before calling backward().

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

Unscales gradients in place and checks for inf/nan.

Returns true if all gradients are finite, false if any overflow.

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

Checks a slice of gradients for inf/nan without modifying them.

pub fn found_inf(&self) -> bool

Returns whether inf/nan was found in the last unscale operation.

pub fn set_found_inf(&mut self, found: bool)

Marks that inf was found (for external gradient checking).

pub fn update(&mut self)

Updates the scale factor based on overflow history.

Call this after each optimizer step:

• If overflow was detected, scale is reduced by backoff_factor
• If no overflow for growth_interval steps, scale is increased by growth_factor

pub fn state_dict(&self) -> GradScalerState

Returns the current state for checkpointing.

pub fn load_state_dict(&mut self, state: GradScalerState)

Loads state from a checkpoint.

Trait Implementations

impl Clone for GradScaler

fn clone(&self) -> GradScaler

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

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.