diff --git a/tensorboard/plugins/interactive_inference/witwidget/BUILD b/tensorboard/plugins/interactive_inference/witwidget/BUILD
index 43932ed6ca..48164a29fd 100644
--- a/tensorboard/plugins/interactive_inference/witwidget/BUILD
+++ b/tensorboard/plugins/interactive_inference/witwidget/BUILD
@@ -9,6 +9,12 @@ py_library(
         "notebook/**/*.py",
     ]),
     srcs_version = "PY2AND3",
+    deps = [
+        "//tensorboard:expect_tensorflow_installed",
+        "//tensorboard/plugins/interactive_inference/utils:inference_utils",
+        "@com_google_protobuf//:protobuf_python",
+        "@org_pythonhosted_six",
+    ]
 )
 
 py_library(
diff --git a/tensorboard/plugins/interactive_inference/witwidget/notebook/base.py b/tensorboard/plugins/interactive_inference/witwidget/notebook/base.py
index e4a9c16d21..d83551907c 100644
--- a/tensorboard/plugins/interactive_inference/witwidget/notebook/base.py
+++ b/tensorboard/plugins/interactive_inference/witwidget/notebook/base.py
@@ -16,10 +16,10 @@
 import json
 import googleapiclient.discovery
 import tensorflow as tf
-import sys
 from IPython import display
 from google.protobuf import json_format
 from numbers import Number
+from six import ensure_str
 from tensorboard.plugins.interactive_inference.utils import inference_utils
 
 # Constants used in mutant inference generation.
@@ -228,10 +228,8 @@ def create_sprite(self):
       example_strings = [
         self.json_to_proto(ex).SerializeToString()
         for ex in self.examples]
-      encoded = base64.b64encode(
-        inference_utils.create_sprite_image(example_strings))
-      if sys.version_info >= (3, 0):
-        encoded = encoded.decode('utf-8')
+      encoded = ensure_str(base64.b64encode(
+        inference_utils.create_sprite_image(example_strings)))
       return 'data:image/png;base64,{}'.format(encoded)
     else:
       return None
@@ -264,7 +262,8 @@ def _json_from_tf_examples(self, tf_examples):
           elif ex.features.feature[feat].HasField('float_list'):
             json_ex[feat_idx] = ex.features.feature[feat].float_list.value[0]
           else:
-            json_ex[feat_idx] = ex.features.feature[feat].bytes_list.value[0]
+            json_ex[feat_idx] = ensure_str(
+              ex.features.feature[feat].bytes_list.value[0])
       else:
         json_ex = {}
         for feat in ex.features.feature:
@@ -275,7 +274,8 @@ def _json_from_tf_examples(self, tf_examples):
           elif ex.features.feature[feat].HasField('float_list'):
             json_ex[feat] = ex.features.feature[feat].float_list.value[0]
           else:
-            json_ex[feat] = ex.features.feature[feat].bytes_list.value[0]
+            json_ex[feat] = ensure_str(
+              ex.features.feature[feat].bytes_list.value[0])
       json_exs.append(json_ex)
     return json_exs