diff --git a/core/leras/nn.py b/core/leras/nn.py
index c4420e1..ef5c2c9 100644
--- a/core/leras/nn.py
+++ b/core/leras/nn.py
@@ -76,15 +76,26 @@ class nn():
             if first_run:
                 io.log_info("Caching GPU kernels...")
 
-            #import tensorflow as tf
-            import tensorflow.compat.v1 as tf
+            import tensorflow
             
+            tf_version = getattr(tensorflow,'VERSION', None)
+            if tf_version is None:
+                tf_version = tensorflow.version.GIT_VERSION
+                if tf_version[0] == 'v':
+                    tf_version = tf_version[1:]
+                
+            if tf_version[0] == '2':
+                tf = tensorflow.compat.v1
+            else:
+                tf = tensorflow
+
             import logging
             # Disable tensorflow warnings
             tf_logger = logging.getLogger('tensorflow')
             tf_logger.setLevel(logging.ERROR)
             
-            tf.disable_v2_behavior()
+            if tf_version[0] == '2':
+                tf.disable_v2_behavior()
             nn.tf = tf
 
             # Initialize framework