Add override FW version

Author: shantanutrip

src/sagemaker/estimator.py

@@ -3547,6 +3547,7 @@ def __init__(
         self.checkpoint_s3_uri = checkpoint_s3_uri
         self.checkpoint_local_path = checkpoint_local_path
         self.enable_sagemaker_metrics = enable_sagemaker_metrics
+        self.override_fw_version = None
 
     def _prepare_for_training(self, job_name=None):
         """Set hyperparameters needed for training. This method will also validate ``source_dir``.
@@ -3731,7 +3732,7 @@ def training_image_uri(self, region=None):
         return image_uris.get_training_image_uri(
             region=region or self.sagemaker_session.boto_region_name,
             framework=self._framework_name,
-            framework_version=self.framework_version,  # pylint: disable=no-member
+            framework_version=self.override_fw_version or self.framework_version,  # pylint: disable=no-member
             py_version=self.py_version,  # pylint: disable=no-member
             image_uri=self.image_uri,
             distribution=getattr(self, "distribution", None),

src/sagemaker/tensorflow/estimator.py

@@ -177,6 +177,7 @@ def __init__(
                 fw.python_deprecation_warning(self._framework_name, defaults.LATEST_PY2_VERSION)
             )
         self.framework_version = framework_version
+        self.override_fw_version = "2.16"
         # TF training and inference versions do not have a one-to-one connection. This mismatch
         # is accommodated by the underlying dictionary. The key of dictionary relates to the inference
         # version and the value relates to training version.