diff --git a/polaris/ocean/tests/baroclinic_channel/viz.py b/polaris/ocean/tests/baroclinic_channel/viz.py
index ece19b037..5e361b768 100644
--- a/polaris/ocean/tests/baroclinic_channel/viz.py
+++ b/polaris/ocean/tests/baroclinic_channel/viz.py
@@ -47,5 +47,5 @@ def run(self):
         plot_horiz_field(config, ds, ds_mesh, 'normalVelocity',
                          'final_normalVelocity.png',
                          t_index=t_index,
-                         vmin=-1 * max_velocity, vmax=max_velocity,
-                         cmap='cmo.balance', show_cell_edges=True)
+                         vmin=-max_velocity, vmax=max_velocity,
+                         cmap='cmo.balance', show_patch_edges=True)
diff --git a/polaris/viz/__init__.py b/polaris/viz/__init__.py
index d745e6d74..2479d0094 100644
--- a/polaris/viz/__init__.py
+++ b/polaris/viz/__init__.py
@@ -82,11 +82,11 @@ def plot_horiz_field(config, ds, ds_mesh, field_name, out_file_name,
     if 'Time' in field.dims and t_index is None:
         t_index = 0
     if t_index is not None:
-        field = field[t_index, :]
+        field = field.isel(Time=t_index)
     if 'nVertLevels' in field.dims and z_index is None:
         z_index = 0
     if z_index is not None:
-        field = field[:, z_index]
+        field = field.isel(nVertLevels=z_index)
 
     if 'nCells' in field.dims:
         ocean_mask = ds_mesh.maxLevelCell - 1 >= 0