diff --git a/src/physics/cam/nn_convection_flux.F90 b/src/physics/cam/nn_convection_flux.F90
index ffa7393c92..b35d4574d4 100644
--- a/src/physics/cam/nn_convection_flux.F90
+++ b/src/physics/cam/nn_convection_flux.F90
@@ -71,9 +71,8 @@ subroutine nn_convection_flux(tabs_i, q_i, &
         ! ---------------------
         ! Fields from beginning of time step used as NN inputs
         ! ---------------------
-        != unit s :: tabs_i
         real(8), intent(in) :: tabs_i(:, :)
-            !! Temperature
+            !! Temperature indices
 
         != unit 1 :: q_i
         real(8), intent(in) :: q_i(:, :)
@@ -130,9 +129,9 @@ subroutine nn_convection_flux(tabs_i, q_i, &
                                                              t_delta_sed, q_delta_sed
             !! delta values of t and q generated by the NN
 
-        != unit kg / m**2 :: precsfc
+        != unit (kg / m**2) :: precsfc
         real(8), intent(out), dimension(:)   :: precsfc
-            !! Surface precipitation rate due to autoconversion and sedimentation
+            !! Surface precipitation due to autoconversion and sedimentation
             !! Can be expressed in m by dividing by density of water 1000 kg / m**3
 
         ! -----------------------------------
@@ -149,6 +148,9 @@ subroutine nn_convection_flux(tabs_i, q_i, &
 
         ! Other variables
         real(8),   dimension(nrf) :: omp, fac
+
+        != unit s m**3 / kg  :: irhoadz
+        != unit s m**2 / kg  :: irhoadzdz
         real(8),   dimension(size(tabs_i, 2)) :: rsat, irhoadz, irhoadzdz
 
         ! -----------------------------------
@@ -159,8 +161,13 @@ subroutine nn_convection_flux(tabs_i, q_i, &
         real(4), dimension(n_outputs) :: outputs
             !! vector of output features from the NN
         ! NN outputs
+
+        != unit kg / (s m**2) :: q_sed_flux
+        ! mass flux of moisture
+        real(8),   dimension(nrf) :: q_sed_flux
+        != unit j / (kg * s) :: t_rad_rest_tend
         real(8),   dimension(nrf) :: t_flux_adv, q_flux_adv, q_tend_auto, &
-                                  q_sed_flux, t_rad_rest_tend
+                                  t_rad_rest_tend
 
         ncol = size(tabs_i, 1)
         nzm = size(tabs_i, 2)
@@ -345,12 +352,18 @@ subroutine nn_convection_flux(tabs_i, q_i, &
             ! Calculate surface precipitation
             ! Combination of sedimentation at surface, and autoconversion in the column
             ! Apply sedimenting flux at surface to get rho*dq term
+
+            ! <temeporary-unit-change>
+            ! unit(kg / m**3) :: precfsc
+            ! NB: precsfc should be all 0 here (and thus of any unit)
             precsfc(i) = precsfc(i) - q_sed_flux(1) * irhoadzdz(1) * rho(1) * adz(1) !!  *dtn/dz
             ! Loop up column for all autoconverted precipitation
             do k=1,nrf
                 precsfc(i) = precsfc(i) - q_delta_auto(i,k) * rho(k) * adz(k)
             end do
             precsfc(i) = precsfc(i) * dz
+            ! unit(kg / m**2) :: precsfc (as specified at the start of the subroutine)
+            ! </temporary-unit-change>
 
             ! As a final check enforce q must be >= 0.0
             do k = 1,nrf