Skip to content

Commit e73b458

Browse files
committed
Add possiblity to accomodate for training/inference version mismatch
1 parent 50fb218 commit e73b458

File tree

1 file changed

+9
-1
lines changed

1 file changed

+9
-1
lines changed

src/sagemaker/tensorflow/model.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,14 @@ def __init__(
185185
)
186186
self.framework_version = framework_version
187187

188+
## Inference framework version is being introduced to accomodate the mismatch between tensorflow and tensorflow serving releases,
189+
## wherein the TF and TFS might have different patch versions, but end up hosting the model of same TF version. For eg., the
190+
## upstream TFS-2.12.0 release was a bad release and hence a new TFS-2.12.1 release was made to host models from TF-2.12.0.
191+
training_inference_version_mismatch_dict = {"2.12.0": "2.12.1"}
192+
self.inference_framework_version = training_inference_version_mismatch_dict.get(
193+
framework_version, framework_version
194+
)
195+
188196
super(TensorFlowModel, self).__init__(
189197
model_data=model_data,
190198
role=role,
@@ -457,7 +465,7 @@ def _get_image_uri(
457465
return image_uris.retrieve(
458466
self._framework_name,
459467
region_name or self.sagemaker_session.boto_region_name,
460-
version=self.framework_version,
468+
version=self.inference_framework_version,
461469
instance_type=instance_type,
462470
accelerator_type=accelerator_type,
463471
image_scope="inference",

0 commit comments

Comments
 (0)