import tensorflow as tf

                train_op = opt.apply_gradients(grads_and_vars, global_step)

        # Validation
        '''
        if params.validation and params.references[0]:
            files = [params.validation] + list(params.references)
            eval_inputs = files
            eval_input_fn = dataset.get_evaluation_input
        else:
            print("Don't evaluate")
            eval_input_fn = None
        '''
        # Add hooks
        train_hooks = [
            tf.train.StopAtStepHook(last_step=params.train_steps),
            tf.train.NanTensorHook(loss),  # Monitors the loss tensor and stops training if loss is NaN
            tf.train.LoggingTensorHook(
                {
                    "step": global_step,
                    "loss": loss,
                    "chars": tf.shape(features["chars"]),
                    "source": tf.shape(features["source"]),
                    #"bert": tf.shape(features["bert"]),
                    "lr": learning_rate
                },
                every_n_iter=1
            ),
            tf.train.CheckpointSaverHook(
