add reqs

Signed-off-by: Ryan Nazareth <[email protected]>
flyteorg · Oct 19, 2022 · 1a99a2f · 1a99a2f
1 parent fa1c100
commit 1a99a2f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/dev-requirements.in b/dev-requirements.in
@@ -13,3 +13,4 @@ google-cloud-bigquery
 google-cloud-bigquery-storage
 IPython
 torch
+tensorflow
diff --git a/flytekit/extras/keras/__init__.py b/flytekit/extras/keras/__init__.py
@@ -13,7 +13,7 @@
 # that have soft dependencies
 try:
     # isolate the exception to the keras import
-    import keras
+    from tensorflow import keras
 
     _keras_installed = True
 except (ImportError, OSError):