code: def _test_rosenbrock(self, constructor, automode=True):
      # automode shows wether we need to update SWA params manually
  
      params = torch.tensor([1.5, 1.5], requires_grad=True)
      optimizer = constructor([params])
  
      solution = torch.tensor([1.0, 1.0])
      initial_dist = params.data.dist(solution)
  
      def eval():
          # SWA
          optimizer.zero_grad()
          loss = rosenbrock(params)
          loss.backward()
          # loss.backward() will give **slightly** different
          # gradients, than drosenbtock, because of a different ordering
          # of floating point operations. In most cases it doesn't matter,
          # but some optimizers are so sensitive that they can temporarily
          # diverge up to 1e-4, just to converge again. This makes the
          # comparison more stable.
          params.grad.data.copy_(drosenbrock(params.data))
          return loss
  
      for _ in range(2000):
          optimizer.step(eval)
          if not automode:
              optimizer.update_swa()
      optimizer.finalize()
  
      self.assertLessEqual(params.data.dist(solution), initial_dist)
  
imports: import contextlib
  import functools
  import re
  import warnings
  from unittest import TestCase
  import torch
  from pytext.optimizer import Lamb, StochasticWeightAveraging
  from torch import nn, optim, sparse
  from torch.autograd import Variable
  from torch.utils import data
label_desc: Use image node_id 2 for calling the TestSWA obj's underlying member method code with example usage: obj._test_rosenbrock(constructor, automode) and returns: loss