code: def step(self) -> None:
      self.t.assign(self.t + 1)
      for i, t in enumerate(self.params):
          assert t.grad is not None
          self.m[i].assign(
              self.b1 * self.m[i] + (1.0 - self.b1) * t.grad
          )
          self.v[i].assign(
              self.b2 * self.v[i] + (1.0 - self.b2) * (t.grad * t.grad)
          )
          m_hat = self.m[i] / (1.0 - self.b1**self.t)
          v_hat = self.v[i] / (1.0 - self.b2**self.t)
          up = (
              m_hat / (v_hat.sqrt() + self.eps)
          ) + self.wd * t.detach()
          if not self.adam:
              r1 = t.detach().square().sum().sqrt()
              r2 = up.square().sum().sqrt()
              r = Tensor.where(
                  r1 > 0, Tensor.where(r2 > 0, r1 / r2, 1.0), 1.0
              )
          else:
              r = 1.0
          t.assign(t.detach() - self.lr * r * up)
      self.realize([self.t] + self.m + self.v)
  
file_path: tinygrad/tinygrad/nn/optim.py
imports: from typing import List
  from tinygrad.helpers import dedup, getenv
  from tinygrad.tensor import Tensor
label_desc: Use image node_id 2 for calling the LAMB obj's underlying member method code with example usage: obj.step() without return types