luz_callback_mixed_precision: Automatic Mixed Precision callback

In mlverse/torchlight: Higher Level 'API' for 'torch'

Automatic Mixed Precision callback

Description

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

See Also

Other luz_callbacks: luz_callback_auto_resume(), luz_callback_csv_logger(), luz_callback_early_stopping(), luz_callback_interrupt(), luz_callback_keep_best_model(), luz_callback_lr_scheduler(), luz_callback_metrics(), luz_callback_mixup(), luz_callback_model_checkpoint(), luz_callback_profile(), luz_callback_progress(), luz_callback_resume_from_checkpoint(), luz_callback_train_valid(), luz_callback()

mlverse/torchlight documentation built on Sept. 19, 2024, 11:22 p.m.