Pass sagemaker_session to estimator

aws · Aug 27, 2018 · 0fd5919 · 0fd5919
1 parent ed4ce41
commit 0fd5919
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/integ/test_tf_keras.py b/tests/integ/test_tf_keras.py
@@ -30,7 +30,7 @@ def test_keras(sagemaker_session, tf_full_version):
     with timeout(minutes=45):
         estimator = TensorFlow(entry_point='keras_cnn_cifar_10.py',
                                source_dir=script_path,
-                               role='SageMakerRole',
+                               role='SageMakerRole', sagemaker_session=sagemaker_session,
                                hyperparameters={'learning_rate': 1e-4, 'decay': 1e-6},
                                training_steps=500, evaluation_steps=5,
                                train_instance_count=1, train_instance_type='ml.c4.xlarge',