import tensorflow as tf

        tf.equal(phase, 'train'),
        datasets.train.get_next,
        datasets.test.get_next)
    if not isinstance(data, dict):
        data = {'data': data}
    if 'length' not in data:
        example = data[list(data.keys())[0]]
        data['length'] = (
                tf.zeros((tf.shape(example)[0],), tf.int32) + tf.shape(example)[1])
    return data


def train(model_fn, datasets, logdir, config):
    """Train a model on a datasets.

    The model function receives the following arguments: data batch, trainer
