diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py index 86a6875721..bf6fcaa376 100644 --- a/src/sagemaker/serve/builder/jumpstart_builder.py +++ b/src/sagemaker/serve/builder/jumpstart_builder.py @@ -157,6 +157,7 @@ def _create_pre_trained_js_model(self) -> Type[Model]: vpc_config=self.vpc_config, sagemaker_session=self.sagemaker_session, name=self.name, + instance_type=self.instance_type, ) self._original_deploy = pysdk_model.deploy diff --git a/tests/integ/sagemaker/serve/test_schema_builder.py b/tests/integ/sagemaker/serve/test_schema_builder.py index 1a2bbe2355..6d3e8281d5 100644 --- a/tests/integ/sagemaker/serve/test_schema_builder.py +++ b/tests/integ/sagemaker/serve/test_schema_builder.py @@ -34,7 +34,9 @@ def test_model_builder_happy_path_with_only_model_id_text_generation(sagemaker_session): model_builder = ModelBuilder( - model="HuggingFaceH4/zephyr-7b-beta", sagemaker_session=sagemaker_session + model="HuggingFaceH4/zephyr-7b-beta", + sagemaker_session=sagemaker_session, + instance_type=None, ) model = model_builder.build(sagemaker_session=sagemaker_session)