diff --git a/jdaviz/core/helpers.py b/jdaviz/core/helpers.py
index e4e4c377b0..a1fb35d954 100644
--- a/jdaviz/core/helpers.py
+++ b/jdaviz/core/helpers.py
@@ -470,7 +470,7 @@ def _handle_display_units(self, data, use_display_units=True):
                 spectral_unit = self.app._get_display_unit('spectral')
                 if not spectral_unit:
                     return data
-                flux_unit = self.app._get_display_unit('flux')
+                y_unit = self.app._get_display_unit('spectral_y')
                 # TODO: any other attributes (meta, wcs, etc)?
                 # TODO: implement uncertainty.to upstream
                 uncertainty = data.uncertainty
@@ -488,18 +488,18 @@ def _handle_display_units(self, data, use_display_units=True):
                         new_uncert = uncertainty
                     if ('_pixel_scale_factor' in data.meta):
                         new_uncert_converted = flux_conversion(data, new_uncert.quantity.value,
-                                                               new_uncert.unit, flux_unit)
-                        new_uncert = StdDevUncertainty(new_uncert_converted, unit=flux_unit)
+                                                               new_uncert.unit, y_unit)
+                        new_uncert = StdDevUncertainty(new_uncert_converted, unit=y_unit)
                     else:
                         new_uncert = StdDevUncertainty(new_uncert, unit=data.flux.unit)
 
                 else:
                     new_uncert = None
                 if ('_pixel_scale_factor' in data.meta):
-                    new_flux = flux_conversion(data, data.flux.value, data.flux.unit,
-                                               flux_unit) * u.Unit(flux_unit)
+                    new_y = flux_conversion(data, data.flux.value, data.flux.unit,
+                                            y_unit) * u.Unit(y_unit)
                 else:
-                    new_flux = flux_conversion(data, data.flux.value, data.flux.unit,
+                    new_y = flux_conversion(data, data.flux.value, data.flux.unit,
                                                data.flux.unit) * u.Unit(data.flux.unit)
                 new_spec = (spectral_axis_conversion(data.spectral_axis.value,
                                                      data.spectral_axis.unit,
@@ -507,7 +507,7 @@ def _handle_display_units(self, data, use_display_units=True):
                             * u.Unit(spectral_unit))
 
                 data = Spectrum1D(spectral_axis=new_spec,
-                                  flux=new_flux,
+                                  flux=new_y,
                                   uncertainty=new_uncert,
                                   mask=data.mask)
             else:  # pragma: nocover
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index a5f1a775ed..bf196988ba 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -2955,17 +2955,6 @@ def _get_continuum(self, dataset, spectral_subset, update_marks=False, per_pixel
                 raise ValueError("per-pixel only supported for cubeviz")
             full_spectrum = self.app._jdaviz_helper.get_data(self.dataset.selected,
                                                              use_display_units=True)
-            # TODO: Something like the following code may be needed to get continuum
-            #  with display units working
-            # temp_spec = self.app._jdaviz_helper.get_data(self.dataset.selected,
-            #                                              use_display_units=True)
-            # flux_values = np.sum(np.ones_like(temp_spec.flux.value), axis=(0, 1))
-            # pix_scale = self.dataset.selected_dc_item.meta.get('PIXAR_SR', 1.0)
-            # pix_scale_factor = (flux_values * pix_scale)
-            # temp_spec.meta['_pixel_scale_factor'] = pix_scale_factor
-            # full_spectrum = self._specviz_helper._handle_display_units(temp_spec,
-            #                                                            use_display_units=True)
-
         else:
             full_spectrum = dataset.get_selected_spectrum(use_display_units=True)