import tensorflow as tf

    else:
        config = tf.estimator.RunConfig(
            model_dir=args.model_dir,
            save_checkpoints_steps=10,
            save_summary_steps=10)

        estimator = tf.estimator.Estimator(
            model_fn,
            config=config,
            params=params)

    estimator.train(train_input_fn, steps=100)
