update

Author: Qingzi-Lan

src/sagemaker/model.py

@@ -626,8 +626,11 @@ def _compilation_job_config(
         ):
             input_model_config["FrameworkVersion"] = utils.get_short_version(framework_version)
 
-        if framework.lower() == "tensorflow":
-            input_model_config["FrameworkVersion"] = utils.get_short_version(framework_version)
+        if (
+            framework.lower() == "tensorflow"
+            and framework_version is not None
+        ):
+            input_model_config["FrameworkVersion"] = framework_version
 
         role = self.sagemaker_session.expand_role(role)
         output_model_config = {

tests/unit/sagemaker/model/test_neo.py

@@ -103,7 +103,6 @@ def test_compile_model_linux_arm64_nvidia(sagemaker_session):
         output_path="s3://output",
         role="role",
         framework="tensorflow",
-        framework_version="1.15.3",
         job_name="compile-model",
         target_platform_os="LINUX",
         target_platform_arch="ARM64",
@@ -124,7 +123,6 @@ def test_compile_model_android_armv7(sagemaker_session):
         output_path="s3://output",
         role="role",
         framework="tensorflow",
-        framework_version="1.15.3",
         job_name="compile-model",
         target_platform_os="ANDROID",
         target_platform_arch="ARM_EABI",