code: def lamb_step(
      p, lr, mom, step, sqr_mom, grad_avg, sqr_avg, eps, **kwargs
  ):
      "Step for LAMB with `lr` on `p`"
      debias1 = debias(mom, 1 - mom, step)
      debias2 = debias(sqr_mom, 1 - sqr_mom, step)
      r1 = p.data.pow(2).mean().sqrt()
      step = (grad_avg / debias1) / ((sqr_avg / debias2).sqrt() + eps)
      r2 = step.pow(2).mean().sqrt()
      q = 1 if r1 == 0 or r2 == 0 else min(r1 / r2, 10)
      p.data.add_(step, alpha=-lr * q)
  
imports: from __future__ import annotations
  from .torch_basics import *
label_desc: Use image node_id 22 for calling a global function with example usage: lamb_step(p, lr, mom, step, sqr_mom, grad_avg, sqr_avg, eps) without return types