diff --git a/dacapo/experiments/tasks/post_processors/watershed_post_processor.py b/dacapo/experiments/tasks/post_processors/watershed_post_processor.py
index 3f658d9ef..b57f07f42 100644
--- a/dacapo/experiments/tasks/post_processors/watershed_post_processor.py
+++ b/dacapo/experiments/tasks/post_processors/watershed_post_processor.py
@@ -132,7 +132,7 @@ def process(
             self.prediction_array_identifier.dataset,
         )
 
-        data = to_ndarray(input_array, output_array.roi)
+        data = to_ndarray(input_array, output_array.roi).astype(float)
         segmentation = mws.agglom(
             data - parameters.bias, offsets=self.offsets, randomized_strides=True
         )
diff --git a/dacapo/predict_local.py b/dacapo/predict_local.py
index 5a49f2c87..76858012b 100644
--- a/dacapo/predict_local.py
+++ b/dacapo/predict_local.py
@@ -57,7 +57,7 @@ def predict(
         output_roi,
         num_channels,
         output_voxel_size,
-        np.uint8,
+        np.float32,
     )
 
     logger.info("Total input ROI: %s, output ROI: %s", input_size, output_roi)
@@ -82,9 +82,6 @@ def predict_fn(block):
                 .cpu()
                 .numpy()[0]
             )
-            predictions = (predictions + 1) * 255.0 / 2.0
-            predictions[predictions > 254] = 0
-            predictions = np.round(predictions).astype(np.uint8)
 
             save_ndarray(predictions, block.write_roi, result_dataset)
             # result_dataset[block.write_roi] = predictions
diff --git a/docs/source/notebooks/minimal_tutorial.py b/docs/source/notebooks/minimal_tutorial.py
index 411612f4d..cf44086a9 100644
--- a/docs/source/notebooks/minimal_tutorial.py
+++ b/docs/source/notebooks/minimal_tutorial.py
@@ -81,6 +81,7 @@
 # import random
 
 import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap
 import numpy as np
 from funlib.geometry import Coordinate, Roi
 from funlib.persistence import prepare_ds
@@ -134,29 +135,18 @@
 )
 labels_array[labels_array.roi] = label(mask_array.to_ndarray(mask_array.roi))[0]
 
-# Generate affinity graph
-affs_array = prepare_ds(
-    "cells3d.zarr",
-    "affs",
-    Roi((0, 0, 0), cell_data.shape[1:]) * voxel_size,
-    voxel_size=voxel_size,
-    num_channels=3,
-    dtype=np.uint8,
-)
-affs_array[affs_array.roi] = (
-    seg_to_affgraph(
-        labels_array.to_ndarray(labels_array.roi),
-        neighborhood=[Coordinate(1, 0, 0), Coordinate(0, 1, 0), Coordinate(0, 0, 1)],
-    )
-    * 255
-)
 print("Data saved to cells3d.zarr")
 
 
+# Create a custom label color map for showing instances
+np.random.seed(1)
+colors = [[0, 0, 0]] + [list(np.random.choice(range(256), size=3)) for _ in range(254)]
+label_cmap = ListedColormap(colors)
+
 # %% [markdown]
 # Here we show a slice of the raw data:
 # %%
-plt.imshow(cell_array.data[30])
+# plt.imshow(cell_array.data[30])
 
 # %% [markdown]
 # ## Datasplit
@@ -177,14 +167,14 @@
         raw_container="cells3d.zarr",
         raw_dataset="raw",
         gt_container="cells3d.zarr",
-        gt_dataset="mask",
+        gt_dataset="labels",
     ),
     DatasetSpec(
         dataset_type="val",
         raw_container="cells3d.zarr",
         raw_dataset="raw",
         gt_container="cells3d.zarr",
-        gt_dataset="mask",
+        gt_dataset="labels",
     ),
 ]
 
@@ -229,7 +219,7 @@
 # an example affinities task configuration
 affs_task_config = AffinitiesTaskConfig(
     name="example_affs",
-    neighborhood=[(0, 1, 0), (0, 0, 1)],
+    neighborhood=[(1, 0, 0), (0, 1, 0), (0, 0, 1)],
 )
 # config_store.delete_task_config(dist_task_config.name)
 config_store.store_task_config(affs_task_config)
@@ -249,8 +239,8 @@
 # all with 1s in z meaning no downsampling or convolving in the z direction.
 architecture_config = CNNectomeUNetConfig(
     name="example_unet",
-    input_shape=(2, 64, 64),
-    eval_shape_increase=(7, 0, 0),
+    input_shape=(2, 132, 132),
+    eval_shape_increase=(8, 32, 32),
     fmaps_in=1,
     num_fmaps=8,
     fmaps_out=8,
@@ -259,7 +249,7 @@
     kernel_size_down=[[(1, 3, 3)] * 2] * 3,
     kernel_size_up=[[(1, 3, 3)] * 2] * 2,
     constant_upsample=True,
-    padding="same",
+    padding="valid",
 )
 config_store.store_architecture_config(architecture_config)
 
@@ -401,17 +391,18 @@
     )[0]
     pred_path = f"/Users/pattonw/dacapo/example_run/validation.zarr/{validation_it}/ds_{dataset}/prediction"
     out_path = f"/Users/pattonw/dacapo/example_run/validation.zarr/{validation_it}/ds_{dataset}/output/WatershedPostProcessorParameters(id=2, bias=0.5, context=(32, 32, 32))"
-    output = zarr.open(
-        out_path
-    )[:]
+    output = zarr.open(out_path)[:]
     prediction = zarr.open(pred_path)[0]
-    print(raw.shape, gt.shape, output.shape)
     c = (raw.shape[1] - gt.shape[1]) // 2
     if c != 0:
         raw = raw[:, c:-c, c:-c]
     ax[validation - 1, 0].imshow(raw[raw.shape[0] // 2])
-    ax[validation - 1, 1].imshow(gt[gt.shape[0] // 2])
+    ax[validation - 1, 1].imshow(
+        gt[gt.shape[0] // 2], cmap=label_cmap, interpolation="none"
+    )
     ax[validation - 1, 2].imshow(prediction[prediction.shape[0] // 2])
-    ax[validation - 1, 3].imshow(output[output.shape[0] // 2])
+    ax[validation - 1, 3].imshow(
+        output[output.shape[0] // 2], cmap=label_cmap, interpolation="none"
+    )
     ax[validation - 1, 0].set_ylabel(f"Validation {validation_it}")
 plt.show()