Optimizer(  # 5.11 MB state
  step=250,
  # Transforms:
  AdamW(learning_rate=0.0002799, b1=0.9, b2=0.999, eps=1e-08, weight_decay=0.1),
  # State:
  mu=MLP(...),  # 669,706 params, 2.55 MB
  nu=MLP(...),  # 669,706 params, 2.55 MB
)