diff --git a/notebooks/keras_classification.ipynb b/notebooks/keras_classification.ipynb
index f43f7427e73..2bd18093176 100644
--- a/notebooks/keras_classification.ipynb
+++ b/notebooks/keras_classification.ipynb
@@ -24,7 +24,7 @@
     "## Step 1 - Requirements\n",
     "To run this notebook, you will need the following packages already installed:\n",
     "* SparseML and SparseZoo;\n",
-    "* Tensorflow >=2.2, which includes Keras and TensorBoard;\n",
+    "* Tensorflow >=2.1, which includes Keras and TensorBoard;\n",
     "* keras2onnx.\n",
     "\n",
     "You can install any package that is not already present via `pip`."
@@ -60,7 +60,7 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "from tensorflow import keras\n",
+    "from sparseml.keras.utils import keras\n",
     "from sparsezoo.models import Zoo\n",
     "\n",
     "# Root directory for the notebook artifacts\n",
@@ -122,6 +122,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import numpy\n",
+    "\n",
     "# Number of classes\n",
     "num_classes = 10\n",
     "\n",
@@ -132,6 +134,10 @@
     "x_train = x_train.astype('float32') / 255\n",
     "x_test = x_test.astype('float32') / 255\n",
     "\n",
+    "# Add batch dimension (for older TF versions)\n",
+    "x_train = numpy.expand_dims(x_train, -1)\n",
+    "x_test = numpy.expand_dims(x_test, -1)\n",
+    "\n",
     "y_train = keras.utils.to_categorical(y_train, num_classes)\n",
     "y_test = keras.utils.to_categorical(y_test, num_classes)\n",
     "\n",
@@ -418,9 +424,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python (pypi_sparseml)",
+   "display_name": "Python (keras_pruning)",
    "language": "python",
-   "name": "pypi_sparseml"
+   "name": "keras_pruning"
   },
   "language_info": {
    "codemirror_mode": {