From e7a06a23e546fa21e3bf640de7a202eb24d240bc Mon Sep 17 00:00:00 2001 From: Erick Benitez-Ramos Date: Thu, 15 Aug 2024 15:59:26 -0700 Subject: [PATCH] fix: pin tf version in tests --- tests/data/tensorflow_mnist/mnist_v2.py | 2 +- tests/integ/test_tf.py | 2 +- tests/integ/test_tuner.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/data/tensorflow_mnist/mnist_v2.py b/tests/data/tensorflow_mnist/mnist_v2.py index bf1750b386..05467dee49 100644 --- a/tests/data/tensorflow_mnist/mnist_v2.py +++ b/tests/data/tensorflow_mnist/mnist_v2.py @@ -198,7 +198,7 @@ def main(args): if args.current_host == args.hosts[0]: ckpt_manager.save() - net.save("/opt/ml/model/1.keras") + net.save("/opt/ml/model/1") if __name__ == "__main__": diff --git a/tests/integ/test_tf.py b/tests/integ/test_tf.py index b03b0e60ec..d0cb326d44 100644 --- a/tests/integ/test_tf.py +++ b/tests/integ/test_tf.py @@ -278,7 +278,7 @@ def _create_and_fit_estimator(sagemaker_session, tf_version, py_version, instanc instance_count=2, instance_type=instance_type, sagemaker_session=sagemaker_session, - framework_version=tf_version, + framework_version="2.14", py_version=py_version, distribution=PARAMETER_SERVER_DISTRIBUTION, disable_profiler=True, diff --git a/tests/integ/test_tuner.py b/tests/integ/test_tuner.py index 3a41ea0094..0c5b025fc5 100644 --- a/tests/integ/test_tuner.py +++ b/tests/integ/test_tuner.py @@ -701,7 +701,7 @@ def test_tuning_tf( instance_count=1, instance_type=cpu_instance_type, sagemaker_session=sagemaker_session, - framework_version=tensorflow_training_latest_version, + framework_version="2.14", py_version=tensorflow_training_latest_py_version, ) @@ -749,7 +749,7 @@ def test_tuning_tf_vpc_multi( entry_point=script_path, source_dir=resource_path, role="SageMakerRole", - framework_version=tensorflow_training_latest_version, + framework_version="2.14", py_version=tensorflow_training_latest_py_version, instance_count=instance_count, instance_type=instance_type,