Skip to content

Commit 4147f69

Browse files
committed
Change tf_full_vesion fixture
1 parent 7a13af7 commit 4147f69

File tree

2 files changed

+7
-4
lines changed

2 files changed

+7
-4
lines changed

src/sagemaker/tensorflow/estimator.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,10 +177,10 @@ def __init__(
177177
fw.python_deprecation_warning(self._framework_name, defaults.LATEST_PY2_VERSION)
178178
)
179179
self.framework_version = framework_version
180-
if self.framework_version and "2.16" in self.framework_version:
181-
self.override_fw_version = "2.16"
182-
# TODO: Revert
183-
print(f"ABCD123 setting self.override_fw_version to {self.override_fw_version}")
180+
# if self.framework_version and "2.16" in self.framework_version:
181+
# self.override_fw_version = "2.16"
182+
# # TODO: Revert
183+
# print(f"ABCD123 setting self.override_fw_version to {self.override_fw_version}")
184184

185185
self.py_version = py_version
186186
self.instance_type = instance_type

tests/conftest.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
from botocore.config import Config
2424
from packaging.version import Version
25+
from packaging.specifiers import SpecifierSet
2526

2627
from sagemaker import Session, image_uris, utils, get_execution_role
2728
from sagemaker.local import LocalSession
@@ -555,6 +556,8 @@ def tf_full_version(tensorflow_training_latest_version, tensorflow_inference_lat
555556
Fixture exists as such, since TF training and TFS have different latest versions.
556557
Otherwise, this would simply be a single latest version.
557558
"""
559+
if Version(tensorflow_training_latest_version) in SpecifierSet(">=2.16"):
560+
return f"{tensorflow_training_latest_version.major}.{tensorflow_training_latest_version.minor}"
558561
return str(
559562
min(
560563
Version(tensorflow_training_latest_version),

0 commit comments

Comments
 (0)