Skip to content

Commit a0a92bd

Browse files
committed
fix: allow custom image when calling deploy or create_model with estimator
1 parent aceb8e0 commit a0a92bd

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

src/sagemaker/pytorch/estimator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,7 @@ def create_model(
175175
code_location=self.code_location,
176176
py_version=self.py_version,
177177
framework_version=self.framework_version,
178-
image=self.image_name,
178+
image=kwargs["image"] if "image" in kwargs else self.image_name,
179179
model_server_workers=model_server_workers,
180180
sagemaker_session=self.sagemaker_session,
181181
vpc_config=self.get_vpc_config(vpc_config_override),

0 commit comments

Comments
 (0)