diff --git a/efficientnet/__init__.py b/efficientnet/__init__.py
index 38cc885..cb9dc72 100644
--- a/efficientnet/__init__.py
+++ b/efficientnet/__init__.py
@@ -34,7 +34,7 @@ def get_submodules_from_kwargs(kwargs):
 
 
 def inject_keras_modules(func):
-    import keras
+    from tensorflow import keras
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         kwargs['backend'] = keras.backend
@@ -60,7 +60,7 @@ def wrapper(*args, **kwargs):
 
 
 def init_keras_custom_objects():
-    import keras
+    from tensorflow import keras
     from . import model
 
     custom_objects = {