code: def create_batch(dataset, hparams, is_training, batch_size=None):
      """Batch a dataset, optional batch_size override."""
      if not batch_size:
          batch_size = hparams.batch_size
      if hparams.max_expected_train_example_len and is_training:
          dataset = dataset.batch(batch_size, drop_remainder=True)
      else:
          dataset = dataset.padded_batch(
              batch_size,
              padded_shapes=dataset.output_shapes,
              drop_remainder=True,
          )
      return dataset
  
imports: import collections
  import functools
  import io
  import re
  import wave
  import zlib
  import librosa
  from magenta.models.onsets_frames_transcription import audio_transform
  from magenta.models.onsets_frames_transcription import constants
  from magenta.models.onsets_frames_transcription import drum_mappings
  from magenta.models.onsets_frames_transcription import melspec_input
  from note_seq import audio_io
  from note_seq import sequences_lib
  from note_seq.protobuf import music_pb2
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 26 for calling a global function with example usage: create_batch(dataset, hparams, is_training, batch_size) and returns: dataset