diff --git a/IntroductionToPlyto.ipynb b/IntroductionToPlyto.ipynb
index dfa01c7..59ac0bc 100644
--- a/IntroductionToPlyto.ipynb
+++ b/IntroductionToPlyto.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "## Introduction to Plyto with Keras\n",
     "#### Python Machine Learning Visualization Toolkit\n",
-    "This notebook will demonstrate how to use our Training Information Callback function with the PlytoAPI to visualize model accuracy and loss throughout the training process of a machine learning algorithm, as well as a tutorial on how to create your own callback function"
+    "This notebook will demonstrate how to use our Keras Accuracy Loss function with the PlytoAPI to visualize model accuracy and loss throughout the training process of a machine learning algorithm, as well as a tutorial on how to create your own callback function"
    ]
   },
   {
@@ -18,11 +18,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using TensorFlow backend.\n"
+     ]
+    }
+   ],
    "source": [
-    "from plyto import PlytoAPI, TrainingInfoCallback"
+    "from plyto import PlytoAPI, KerasAccuracyLossCallback"
    ]
   },
   {
@@ -41,21 +49,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
-   "outputs": [
-    {
-     "ename": "NameError",
-     "evalue": "name 'PlytoAPI' is not defined",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "\u001b[0;32m<ipython-input-1-ab255c47190c>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     37\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[0;31m# pass the spec into PlytoAPI\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 39\u001b[0;31m \u001b[0mplytoInstance\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPlytoAPI\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mspec\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     40\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     41\u001b[0m \u001b[0;31m# pass the plyto instance into the callback function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'PlytoAPI' is not defined"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# an array of altair specs with one plot of samples versus accuracy\n",
     "spec = [\n",
@@ -98,7 +94,7 @@
     "plytoInstance = PlytoAPI(spec)\n",
     "\n",
     "# pass the plyto instance into the callback function\n",
-    "callback = TrainingInfoCallback(plytoInstance)"
+    "callback = KerasAccuracyLossCallback(plytoInstance)"
    ]
   },
   {
@@ -111,7 +107,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -221,7 +217,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -291,13 +287,13 @@
     "            self.plyto.send_data()\n",
     "\n",
     "        \n",
-    "basicPlatoInstance = PlytoAPI({})\n",
+    "basicPlatoInstance = PlytoAPI([])\n",
     "basicCallback = Basic(basicPlatoInstance)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 37,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -307,7 +303,7 @@
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
       "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
-      "\u001b[0;32m<ipython-input-37-db1f4865e5f9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mY_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcallbacks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbasicCallback\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+      "\u001b[0;32m<ipython-input-5-db1f4865e5f9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mY_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcallbacks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbasicCallback\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
       "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, **kwargs)\u001b[0m\n\u001b[1;32m   1035\u001b[0m                                         \u001b[0minitial_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minitial_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1036\u001b[0m                                         \u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1037\u001b[0;31m                                         validation_steps=validation_steps)\n\u001b[0m\u001b[1;32m   1038\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1039\u001b[0m     def evaluate(self, x=None, y=None,\n",
       "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/keras/engine/training_arrays.py\u001b[0m in \u001b[0;36mfit_loop\u001b[0;34m(model, f, ins, out_labels, batch_size, epochs, verbose, callbacks, val_f, val_ins, shuffle, callback_metrics, initial_epoch, steps_per_epoch, validation_steps)\u001b[0m\n\u001b[1;32m    197\u001b[0m                     \u001b[0mins_batch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mins_batch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtoarray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 199\u001b[0;31m                 \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mins_batch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    200\u001b[0m                 \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mto_list\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    201\u001b[0m                 \u001b[0;32mfor\u001b[0m \u001b[0ml\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout_labels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m   2670\u001b[0m                     \u001b[0;34m'In order to feed symbolic tensors to a Keras model '\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2671\u001b[0m                     'in TensorFlow, you need tensorflow 1.8 or higher.')\n\u001b[0;32m-> 2672\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_legacy_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2673\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2674\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
diff --git a/src/StatusItemWidget.tsx b/src/StatusItemWidget.tsx
index 90beeab..3c6dc46 100644
--- a/src/StatusItemWidget.tsx
+++ b/src/StatusItemWidget.tsx
@@ -82,6 +82,24 @@ class StatusItem extends React.Component<IStatusItemProps, IStatusItemState> {
         );
       }
     });
+
+    this.props.tracker.currentWidget.session.kernelChanged.connect(() => {
+      let widget: NotebookPanel | null = this.props.tracker.currentWidget;
+      if (widget) {
+        console.log('new widget. re-registering comm targets');
+
+        this.setState(
+          {
+            kernel: widget.session.kernel as Kernel.IKernel
+          },
+          () => {
+            this.state.kernel.iopubMessage.connect(this.onMessage, this);
+
+            this.state.kernel.registerCommTarget('plyto', (comm, msg) => {});
+          }
+        );
+      }
+    })
   }
 
   onMessage(sender: Kernel.IKernel, msg: KernelMessage.IIOPubMessage) {
diff --git a/src/components/ModelViewer.tsx b/src/components/ModelViewer.tsx
index 4b3afc1..1fcd033 100644
--- a/src/components/ModelViewer.tsx
+++ b/src/components/ModelViewer.tsx
@@ -50,7 +50,7 @@ export class ModelViewer extends React.Component<IModelViewerProps, {}> {
           )}
         </div>
         <div className={graphsStyle}>
-          {this.props.spec.map(spec => {
+          {this.props.spec !== [] && this.props.spec.map(spec => {
             return (
               <div
                 key={spec['name']}