fix: Import modules according to TensorFlow versions

Kohulan · Jul 31, 2023 · 65600ff · 65600ff
1 parent e1b946a
commit 65600ff
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/DECIMER/Efficient_Net_encoder.py b/DECIMER/Efficient_Net_encoder.py
@@ -1,6 +1,8 @@
 # EfficientNet-V2 config
 import tensorflow as tf
-import DECIMER.efficientnetv2 as efficientnetv2
+import efficientnetv2
+from efficientnetv2 import effnetv2_model
+from efficientnetv2 import effnetv2_configs
 
 BATCH_SIZE_DEBUG = 2
 MODEL = "efficientnetv2-m"  # @param

diff --git a/DECIMER/Predictor_usingCheckpoints.py b/DECIMER/Predictor_usingCheckpoints.py
@@ -6,7 +6,11 @@
 import pystow
 from selfies import decoder
 import Transformer_decoder
-import Efficient_Net_encoder
+
+if int(tf.__version__.split(".")[1]) <= 10:
+    import Efficient_Net_encoder
+else:
+    raise ImportError("Please use tensorflow 2.10 when working with the checkpoints.")
 import config
 import utils