diff --git a/tasks/imdb_tcn.py b/tasks/imdb_tcn.py
index 9580749..cb5a0ab 100644
--- a/tasks/imdb_tcn.py
+++ b/tasks/imdb_tcn.py
@@ -5,10 +5,18 @@
 Based on: https://github.com/keras-team/keras/blob/master/examples/imdb_bidirectional_lstm.py
 """
 import numpy as np
-from keras import Model, Input
-from keras.datasets import imdb
-from keras.layers import Dense, Dropout, Embedding
-from keras.preprocessing import sequence
+try:
+    from tensorflow.python import keras
+    from tensorflow.python.keras import Model, Input
+    from tensorflow.python.keras.datasets import imdb
+    from tensorflow.python.keras.layers import Dense, Dropout, Embedding
+    from tensorflow.python.keras.preprocessing import sequence
+except:
+    import keras
+    from keras import Model, Input
+    from keras.datasets import imdb
+    from keras.layers import Dense, Dropout, Embedding
+    from keras.preprocessing import sequence
 
 from tcn import TCN
 
diff --git a/tasks/multi_length_sequences.py b/tasks/multi_length_sequences.py
index cb8c94f..6706c29 100644
--- a/tasks/multi_length_sequences.py
+++ b/tasks/multi_length_sequences.py
@@ -1,6 +1,11 @@
 import numpy as np
-from keras.layers import Dense
-from keras.models import Input, Model
+
+try:
+    from tensorflow.python.keras.layers import Dense
+    from tensorflow.python.keras.models import Input, Model
+except:
+    from keras.layers import Dense
+    from keras.models import Input, Model
 
 from tcn import TCN
 
diff --git a/tasks/save_reload_model.py b/tasks/save_reload_model.py
index 68dd541..8907670 100644
--- a/tasks/save_reload_model.py
+++ b/tasks/save_reload_model.py
@@ -1,8 +1,13 @@
 import os
 
 import numpy as np
-from keras import Model, Input
-from keras.layers import Dense, Dropout, Embedding
+
+try:
+    from tensorflow.python.keras import Model, Input
+    from tensorflow.python.keras.layers import Dense, Dropout, Embedding
+except:
+    from keras import Model, Input
+    from keras.layers import Dense, Dropout, Embedding
 
 from tcn import TCN
 
diff --git a/tasks/time_series_forecasting.py b/tasks/time_series_forecasting.py
index 9479d60..a6235da 100644
--- a/tasks/time_series_forecasting.py
+++ b/tasks/time_series_forecasting.py
@@ -1,9 +1,15 @@
 # https://datamarket.com/data/set/22ox/monthly-milk-production-pounds-per-cow-jan-62-dec-75#!ds=22ox&display=line
+import os
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-from keras import Input, Model
-from keras.layers import Dense
+
+try:
+    from tensorflow.python.keras import Input, Model
+    from tensorflow.python.keras.layers import Dense
+except:
+    from keras import Input, Model
+    from keras.layers import Dense
 
 from tcn import TCN
 
@@ -14,7 +20,8 @@
 # - The model is simple.
 ##
 
-milk = pd.read_csv('monthly-milk-production-pounds-p.csv', index_col=0, parse_dates=True)
+filepath = os.path.join(os.path.dirname(__file__), 'monthly-milk-production-pounds-p.csv')
+milk = pd.read_csv(filepath, index_col=0, parse_dates=True)
 
 print(milk.head())
 
diff --git a/tcn/tcn.py b/tcn/tcn.py
index e60df70..54ad011 100644
--- a/tcn/tcn.py
+++ b/tcn/tcn.py
@@ -1,13 +1,25 @@
 from typing import List, Tuple
 
-import keras.backend as K
-import keras.layers
-from keras import optimizers
-from keras.engine.topology import Layer
-from keras.layers import Activation, Lambda
-from keras.layers import Conv1D, SpatialDropout1D
-from keras.layers import Dense, BatchNormalization
-from keras.models import Input, Model
+try:
+    from tensorflow.python import keras
+    from tensorflow.python.keras import backend as k
+    from tensorflow.python.keras import layers
+    from tensorflow.python.keras import optimizers
+    from tensorflow.python.keras.engine import Layer
+    from tensorflow.python.keras.layers import Activation, Lambda
+    from tensorflow.python.keras.layers import Conv1D, SpatialDropout1D
+    from tensorflow.python.keras.layers import Dense, BatchNormalization
+    from tensorflow.python.keras.models import Input, Model
+except:
+    import keras
+    import keras.backend as K
+    from keras import layers
+    from keras import optimizers
+    from keras.engine.topology import Layer
+    from keras.layers import Activation, Lambda
+    from keras.layers import Conv1D, SpatialDropout1D
+    from keras.layers import Dense, BatchNormalization
+    from keras.models import Input, Model
 
 
 def residual_block(x, training, dilation_rate, nb_filters, kernel_size, padding, activation='relu', dropout_rate=0,
@@ -44,7 +56,7 @@ def residual_block(x, training, dilation_rate, nb_filters, kernel_size, padding,
 
     # 1x1 conv to match the shapes (channel dimension).
     prev_x = Conv1D(nb_filters, 1, padding='same')(prev_x)
-    res_x = keras.layers.add([prev_x, x])
+    res_x = layers.add([prev_x, x])
     res_x = Activation(activation)(res_x)
     return res_x, x
 
@@ -140,7 +152,7 @@ def __call__(self, inputs, training=None):
                                              use_batch_norm=self.use_batch_norm)
                 skip_connections.append(skip_out)
         if self.use_skip_connections:
-            x = keras.layers.add(skip_connections)
+            x = layers.add(skip_connections)
         if not self.return_sequences:
             x = Lambda(lambda tt: tt[:, -1, :])(x)
         return x