code: def _build_model(two_heads):
      cls_head = tl.Dense(10)
      if two_heads:
          reg_head = tl.Dense(1)
          heads = tl.Branch(cls_head, reg_head)
      else:
          heads = cls_head
      return tl.Serial(
          tl.Fn("ScaleInput", lambda x: x / 255),
          tl.Flatten(),
          tl.Dense(512),
          tl.Relu(),
          tl.Dense(512),
          tl.Relu(),
          heads,
      )
  
file_path: trax/trax/supervised/mnist_test.py
imports: import io
  import itertools
  from unittest import mock
  from absl.testing import absltest
  from trax import layers
  from trax.data import inputs
  from trax.data import tf_inputs
  from trax.optimizers import adam
  from trax.supervised import training
label_desc: Use image node_id 1 for calling a global function with example usage: _build_model(two_heads) and returns: tl