Function one_cycle_lr 

pub fn one_cycle_lr(max_lr: f32, total_steps: u32, step: u32) -> Result<f32>

One-cycle learning rate at step (0-indexed of total_steps), matching PyTorch’s OneCycleLR at its own defaults (https://docs.pytorch.org/docs/2.14/generated/torch.optim.lr_scheduler.OneCycleLR.html): a cosine warmup from max_lr / div_factor up to max_lr over the first pct_start fraction of total_steps, then a cosine anneal down to max_lr / div_factor / final_div_factor. PyTorch’s own closed form (torch/optim/lr_scheduler.py, _annealing_cos and OneCycleLR.get_lr):

annealing_cos(start, end, pct) = end + (start - end) / 2 * (cos(pi * pct) + 1)
step_size_up   = pct_start * total_steps - 1
step_size_down = (total_steps - 1) - step_size_up
lr(step) = annealing_cos(initial_lr, max_lr, step / step_size_up)                if step <= step_size_up
         = annealing_cos(max_lr, min_lr, (step - step_size_up) / step_size_down) otherwise

Call SGD::set_learning_rate, Adam::set_learning_rate, or AdamW::set_learning_rate with the result once per step, before that step’s update. step must be strictly less than total_steps, matching PyTorch’s own “Tried to step … but the specified number of total steps …” guard against scheduling past the cycle.