code: def setUp(self):
      super().setUp()
      test_utils.ensure_flag("test_tmpdir")
      self._model_fn = functools.partial(
          models.Policy,
          body=lambda mode: tl.Serial(  # pylint: disable=g-long-lambda
              tl.Dense(64), tl.Relu(), tl.Dense(64), tl.Relu()
          ),
      )
  
imports: import functools
  import math
  import os
  import pickle
  from absl.testing import absltest
  import tensorflow
  from trax import layers
  from trax import models
  from trax import optimizers
  from trax import test_utils
  from trax.rl import task
  from trax.rl import training
  from trax.supervised import lr_schedules
label_desc: Use image node_id 1 for calling the TrainingTest obj's underlying member method code with example usage: obj.setUp() without return types