Add version check

optuna · Oct 25, 2024 · d8ef118 · d8ef118
1 parent e7b117c
commit d8ef118
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/tensorflow/tensorflow_estimator_integration.py b/tensorflow/tensorflow_estimator_integration.py
@@ -12,6 +12,7 @@
 """
 
 import shutil
+import sys
 import tempfile
 import urllib
 
@@ -133,4 +134,11 @@ def main():
 
 
 if __name__ == "__main__":
+    if sys.version_info >= (3, 12):
+        raise RuntimeError(
+            "Use `optuna_integration.KerasPruningCallback` instead, "
+            "because TensorFlow recommends migrating from Estimator to Keras APIs. "
+            "For more details, please check the TensorFlow migration guide: "
+            "https://www.tensorflow.org/guide/migrate/migrating_estimator"
+        )
     main()
diff --git a/tensorflow/tensorflow_estimator_simple.py b/tensorflow/tensorflow_estimator_simple.py
@@ -11,6 +11,7 @@
 import shutil
 import tempfile
 import urllib
+import sys
 
 import optuna
 from packaging import version
@@ -123,4 +124,10 @@ def main():
 
 
 if __name__ == "__main__":
+    if sys.version_info >= (3, 12):
+        raise RuntimeError(
+            "Use Keras APIs instead of Estimator API in TensorFlow. "
+            "For more details, please check the TensorFlow migration guide: "
+            "https://www.tensorflow.org/guide/migrate/migrating_estimator"
+        )
     main()