Struct GradScaler 

pub struct GradScaler { /* private fields */ }

GradScaler for mixed precision training.

Scales loss before backward to prevent gradient underflow in float16, then unscales gradients before optimizer step. Dynamically adjusts scale factor based on whether inf/nan gradients are detected.

let mut scaler = GradScaler::new();
let scaled_loss = scaler.scale(&loss)?;
scaled_loss.backward()?;
let stepped = scaler.step(&params, &mut || optim.step())?;
scaler.update();

Implementations

impl GradScaler

pub fn new() -> Self

Create a new GradScaler with default settings.

Initial scale: 2^16 = 65536, growth: 2.0, backoff: 0.5, interval: 2000.

pub fn scale(&self, loss: &Variable) -> Result<Variable>

Scale the loss before backward. Returns loss * scale.

pub fn scale_factor(&self) -> f64

Current scale factor.

pub fn step( &mut self, params: &[Parameter], step_fn: &mut dyn FnMut() -> Result<()>, ) -> Result<bool>

Unscale gradients, check for inf/nan, and step the optimizer.

Returns true if the step was taken (all gradients finite). Returns false if inf/nan detected (optimizer step skipped).

pub fn update(&mut self)

Update the scale factor after each step.

Call this after every step() call, regardless of whether it succeeded.

Trait Implementations

impl Default for GradScaler

fn default() -> Self

Returns the “default value” for a type.

impl Stateful for GradScaler

fn save_state<W: Write>(&self, w: &mut W) -> Result<()>

Serialize optimizer state (lr, momentum buffers, etc.) to a writer.

fn load_state<R: Read>(&mut self, r: &mut R) -> Result<()>

Restore optimizer state from a reader.

fn save_state_file(&self, path: &str) -> Result<()>

Save state to a file. Uses gzip compression if path ends with .gz.

fn load_state_file(&mut self, path: &str) -> Result<()>

Load state from a file. Detects gzip from .gz extension.