Skip to contents

This callback will enable torch::local_autocast() training model forward and during loss computation. It will then disable autocast and scale the loss before backward() and opt$step(). See here for more information.

Usage

luz_callback_mixed_precision(...)

Arguments

...

Passed to torch::cuda_amp_grad_scaler().

Value

A luz_callback