torch.optim.functional.adamw

torch.optim.functional.adamw(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, foreach=None, capturable=False, differentiable=False, fused=None, grad_scale=None, found_inf=None, has_complex=False, *, amsgrad, beta1, beta2, lr, weight_decay, eps, maximize)

Functional API that performs AdamW algorithm computation.

This function updates the provided parameters and optimizer state in place. The caller must initialize and retain optimizer state. Unless intentionally constructing a differentiable update with a supported differentiable=True argument, call this function under torch.no_grad. See Functional optimizer API for the common functional optimizer contract and examples, and AdamW for algorithm details.

Note

With fused=True, CUDA supports FP32 parameters and gradients with BF16 moment buffers. See AdamW with BF16 optimizer state for an example.