code: def test_adam_w(self):
      w = torch.tensor([0.1, -0.2, -0.1], requires_grad=True)
      target = torch.tensor([0.4, 0.2, -0.5])
      criterion = nn.MSELoss()
      # No warmup, constant schedule, no gradient clipping
      optimizer = AdamW(params=[w], lr=2e-1, weight_decay=0.0)
      for _ in range(100):
          loss = criterion(w, target)
          loss.backward()
          optimizer.step()
          w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
          w.grad.zero_()
      self.assertListAlmostEqual(w.tolist(), [0.4, 0.2, -0.5], tol=1e-2)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/tests/optimization/test_optimization.py
imports: import os
  import tempfile
  import unittest
  from transformers import is_torch_available
  from transformers.testing_utils import require_torch
label_desc: Use image node_id 2 for calling the OptimizationTest obj's underlying member method code with example usage: obj.test_adam_w() without return types