Fixed cpu isntance type for the estimator register test (#4811)

aws · Jul 29, 2024 · e09693c · e09693c
1 parent fa68f1f
commit e09693c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/integ/test_byo_estimator.py b/tests/integ/test_byo_estimator.py
@@ -109,7 +109,7 @@ def test_byo_estimator(sagemaker_session, region, cpu_instance_type, training_se
 
 
 @pytest.mark.release
-def test_estimator_register_publish_training_details(sagemaker_session, region):
+def test_estimator_register_publish_training_details(sagemaker_session, region, cpu_instance_type):
 
     bucket = sagemaker_session.default_bucket()
     prefix = "model-card-sample-notebook"
@@ -150,7 +150,7 @@ def test_estimator_register_publish_training_details(sagemaker_session, region):
         container,
         role="SageMakerRole",
         instance_count=1,
-        instance_type="ml.m4.xlarge",
+        instance_type=cpu_instance_type,
         output_path=output_location,
         sagemaker_session=sagemaker_session,
     )