diff --git a/notebooks/wut-predict.ipynb b/notebooks/wut-predict.ipynb
index fbb4633..d2a8eab 100644
--- a/notebooks/wut-predict.ipynb
+++ b/notebooks/wut-predict.ipynb
@@ -34,7 +34,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [],
    "source": [
     "print(\"Start\")"
@@ -65,18 +67,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import keras\n",
-    "from keras import layers\n",
-    "from keras import Sequential\n",
-    "from keras.layers import Activation, Dropout, Flatten, Dense\n",
-    "from keras.preprocessing.image import ImageDataGenerator\n",
-    "from keras.layers import Convolution2D, MaxPooling2D, ZeroPadding2D\n",
-    "from keras import optimizers\n",
-    "from keras.preprocessing import image\n",
-    "from keras.models import load_model\n",
-    "#from keras.preprocessing.image import load_img\n",
-    "#from keras.preprocessing.image import img_to_array\n",
-    "from keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D"
+    "import tensorflow as tf\n",
+    "from tensorflow import keras\n",
+    "from tensorflow.keras import layers\n",
+    "from tensorflow.keras import optimizers\n",
+    "from tensorflow.keras import Sequential\n",
+    "from tensorflow.keras.layers import Activation, Dropout, Flatten, Dense\n",
+    "from tensorflow.keras.layers import Convolution2D, MaxPooling2D, ZeroPadding2D\n",
+    "from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D\n",
+    "from tensorflow.keras.layers import Input, concatenate\n",
+    "from tensorflow.keras.models import load_model\n",
+    "from tensorflow.keras.models import Model\n",
+    "from tensorflow.keras.preprocessing import image\n",
+    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
+    "from tensorflow.keras.preprocessing.image import img_to_array\n",
+    "from tensorflow.keras.preprocessing.image import load_img\n",
+    "from tensorflow.keras.utils import model_to_dot\n",
+    "from tensorflow.keras.utils import plot_model"
    ]
   },
   {
@@ -226,7 +233,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [],
    "source": [
     "sample_test_images, _ = next(test_data_gen)"