code: def model_fn(features, labels, mode, params, config):
      """Builds the acoustic (for Estimator API)."""
      del config
      hparams = params
  
      length = features.length
      spec = features.spec
      is_training = mode == tf_estimator.ModeKeys.TRAIN
  
      (
          frame_logits,
          onset_logits,
          offset_logits,
          velocity_values,
      ) = build_model(spec, length, hparams, is_training)
  
      # Hack to restore the batch dimension, which is lost in some cases.
      def fix_shape(output):
          shape = list(output.shape)
          shape[0] = spec.shape[0]
          output.set_shape(shape)
  
      for output in (
          frame_logits,
          onset_logits,
          offset_logits,
          velocity_values,
      ):
          fix_shape(output)
  
      return estimator_spec_util.get_estimator_spec(
          hparams,
          mode,
          features,
          labels,
          frame_logits,
          onset_logits,
          offset_logits,
          velocity_values,
          offset_network=hparams.offset_network,
      )
  
file_path: magenta/magenta/models/onsets_frames_transcription/model_tpu.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from magenta.contrib import rnn
  from magenta.contrib import training
  from magenta.models.onsets_frames_transcription import constants
  from magenta.models.onsets_frames_transcription import estimator_spec_util
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tf_slim
label_desc: Use image node_id 6 for calling a global function with example usage: model_fn(features, labels, mode, params, config) and returns: estimator_spec_util