Relax TF_KERAS boolean (#6)

CyberZHG · Sep 23, 2019 · 0b3a96f · 0b3a96f
1 parent ea1b9cc
commit 0b3a96f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/keras_lookahead/__init__.py b/keras_lookahead/__init__.py
@@ -1,3 +1,3 @@
 from .optimizers import *
 
-__version__ = '0.7.0'
+__version__ = '0.8.0'
diff --git a/keras_lookahead/backend.py b/keras_lookahead/backend.py
@@ -1,5 +1,8 @@
 import os
 
+from distutils.util import strtobool
+
+
 __all__ = [
     'keras', 'utils', 'activations', 'applications', 'backend', 'datasets', 'engine',
     'layers', 'preprocessing', 'wrappers', 'callbacks', 'constraints', 'initializers',
@@ -9,11 +12,11 @@
 TF_KERAS = False
 EAGER_MODE = False
 
-if os.environ.get('TF_KERAS', '0') != '0':
+if strtobool(os.environ.get('TF_KERAS', '0')):
     import tensorflow as tf
     from tensorflow.python import keras
     TF_KERAS = True
-    if os.environ.get('TF_EAGER', '0') != '0':
+    if strtobool(os.environ.get('TF_EAGER', '0')):
         try:
             tf.enable_eager_execution()
             raise AttributeError()