changes based on cr feedback

Author: icywang86rui

src/sagemaker/utils.py

@@ -123,6 +123,14 @@ def get_config_value(key_path, config):
 
 
 def get_short_version(framework_version):
+    """Return short version in the format of x.x
+
+    Args:
+        framework_version: The version string to be shortened.
+
+    Returns:
+        str: The short version string
+    """
     return '.'.join(framework_version.split('.')[:2])
 
 

tests/unit/test_tf_estimator.py

@@ -743,6 +743,10 @@ def test_legacy_mode_deprecated(sagemaker_session):
                    py_version='py2', script_mode=False)
     assert tf._script_mode_enabled() is True
 
+    tf = _build_tf(sagemaker_session=sagemaker_session, framework_version='1.12',
+                   py_version='py2', script_mode=False)
+    assert tf._script_mode_enabled() is False
+
 
 def test_script_mode_enabled(sagemaker_session):
     tf = _build_tf(sagemaker_session=sagemaker_session, py_version='py3')