diff --git a/generator/ctrl/training_utils/training.py b/generator/ctrl/training_utils/training.py
index 8134076..bc2f13c 100644
--- a/generator/ctrl/training_utils/training.py
+++ b/generator/ctrl/training_utils/training.py
@@ -144,7 +144,7 @@ print(model.summary())
 # this is where the saved model is presented to the code
 # the model directory should have the model checkpoint and
 # a checkpoint file
-run_config = tf.contrib.tpu.RunConfig(
+run_config = tf.estimator.RunConfig(
         model_dir=args.model_dir)