soz (#42)

Co-authored-by: Helen Ngo <[email protected]>
for-ai · Apr 12, 2021 · 79f9e0a · 79f9e0a
1 parent e46dc85
commit 79f9e0a
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/cloud/envs/gcp.py b/cloud/envs/gcp.py
@@ -227,7 +227,11 @@ def __init__(self, instance):
             import tensorflow as tf
             import re
             m = re.search(r'(\d+\.\d+)\.\d+', tf.__version__)
-            self.tf_version = m.group()
+            self.core_tf_version = m.group(1)
+            if float(self.core_tf_version) <= 2.3:
+                self.tf_version = m.group(1)
+            else:
+                self.tf_version = m.group()
             if "dev" in tf.__version__:
                 logging.info("Found Tensorflow nightly version. Using TPU software version: nightly")
                 self.tf_version = "nightly"