diff --git a/NN_module/nn_convection_flux.f90 b/NN_module/nn_convection_flux.f90
index 6ec629e..90de0d6 100644
--- a/NN_module/nn_convection_flux.f90
+++ b/NN_module/nn_convection_flux.f90
@@ -57,12 +57,9 @@ end subroutine nn_convection_flux_init
 
     subroutine nn_convection_flux(tabs_i, q_i, y_in, &
                                   tabs, &
-                                  t_0, q_0, &
+                                  t, q, &
                                   rho, adz, dz, dtn, &
-                                  t_rad_rest_tend, &
-                                  t_delta_adv, q_delta_adv, &
-                                  t_delta_auto, q_delta_auto, &
-                                  t_delta_sed, q_delta_sed, prec_sed)
+                                  prec_sed)
         !! Interface to the neural net that applies physical constraints and reshaping
         !! of variables.
         !! Operates on subcycle of timestep dtn to update t, q, precsfc, and prec_xy
@@ -77,35 +74,29 @@ subroutine nn_convection_flux(tabs_i, q_i, y_in, &
         != unit s :: tabs_i
         real, intent(in) :: tabs_i(:, :)
             !! Temperature
-        
+
         != unit 1 :: q_i
         real, intent(in) :: q_i(:, :)
             !! Non-precipitating water mixing ratio
-        
-        != unit m :: y_in
-        real, intent(in) :: y_in(:)
-            !! Distance of column from equator (proxy for insolation and sfc albedo)
 
         ! ---------------------
         ! Other fields from SAM
         ! ---------------------
+        != unit m :: y_in
+        real, intent(in) :: y_in(:)
+            !! Distance of column from equator (proxy for insolation and sfc albedo)
+
         != unit K :: tabs
         real, intent(in) :: tabs(:, :)
             !! absolute temperature
         
         != unit (J / kg) :: t
-        real, intent(in) :: t_0(:, :)
+        real, intent(inout) :: t(:, :)
             !! Liquid Ice static energy (cp*T + g*z − L(qliq + qice) − Lf*qice)
-        != unit (J / kg) :: t
-        real, allocatable :: t(:, :)
-            !! Copy of t_0 to ensure that t_0 is not modified by this routine for SAM
         
         != unit 1 :: q
-        real, intent(in) :: q_0(:, :)
+        real, intent(inout) :: q(:, :)
             !! total water
-        != unit 1 :: q
-        real, allocatable :: q(:, :)
-            !! Copy of q_0 to ensure that q_0 is not modified by this routine for SAM
 
         ! ---------------------
         ! reference vertical profiles:
@@ -138,14 +129,14 @@ subroutine nn_convection_flux(tabs_i, q_i, y_in, &
         ! -----------------------------------
         != unit (J / kg) :: t_delta_adv, t_delta_auto, t_delta_sed
         != unit 1 :: q_delta_adv, q_delta_auto, q_delta_sed
-        real, intent(out), dimension(:, :) :: t_delta_adv, q_delta_adv, &
-                                              t_delta_auto, q_delta_auto, &
-                                              t_delta_sed, q_delta_sed
-            !! delta values of t and q to be returned by the scheme
+        real, dimension(size(tabs_i, 1), size(tabs_i, 2)) :: t_delta_adv, q_delta_adv, &
+                                                             t_delta_auto, q_delta_auto, &
+                                                             t_delta_sed, q_delta_sed
+            !! delta values of t and q generated by the NN
 
         != unit kg :: t_rad_rest_tend
-        real, intent(out), dimension(:, :) :: t_rad_rest_tend
-            !! tendency of t to be returned by the scheme
+        real, dimension(size(tabs_i, 1), size(tabs_i, 2)) :: t_rad_rest_tend
+            !! tendency of t generated by the NN
         
         != unit kg :: prec_sed
         real, intent(out), dimension(:)   :: prec_sed
@@ -182,11 +173,6 @@ subroutine nn_convection_flux(tabs_i, q_i, y_in, &
         nx = size(tabs_i, 1)
         nzm = size(tabs_i, 2)
 
-        allocate(t(nx, nrf))
-        allocate(q(nx, nrf))
-        q = q_0
-        t = t_0
-
         ! Check that we have initialised all of the variables.
         if (do_init) call error_mesg('NN has not yet been initialised using nn_convection_flux_init.')
 
@@ -360,30 +346,24 @@ subroutine nn_convection_flux(tabs_i, q_i, y_in, &
                 end if
             end do
 
-            ! These final updates are now redundant as t and q updated in interface
-            ! ! Update q and t with sed q and t deltas (dt = -dq*(latent_heat/cp))
-            ! q(i,j,1:nrf) = q(i,j,1:nrf) + q_delta_sed(i,j,1:nrf)
-            ! t(i,j,1:nrf) = t(i,j,1:nrf) - q_delta_sed(i,j,1:nrf)*(fac_fus+fac_cond)
-            t_delta_sed(i,1:nrf) = - q_delta_sed(i,1:nrf)*(fac_fus+fac_cond)
+            ! Update q and t with sed q and t deltas (dt = -dq*(latent_heat/cp))
+            q(i,1:nrf) = q(i,1:nrf) + q_delta_sed(i,1:nrf)
+            t(i,1:nrf) = t(i,1:nrf) - q_delta_sed(i,1:nrf)*(fac_fus+fac_cond)
             !
-            ! ! Apply radiation rest tendency to variables (multiply by dtn to get dt)
-            ! t(i,j,1:nrf) = t(i,j,1:nrf) + t_rad_rest_tend(i,j,1:nrf)*dtn
+            ! Apply radiation rest tendency to variables (multiply by dtn to get dt)
+            t(i,1:nrf) = t(i,1:nrf) + t_rad_rest_tend(i,1:nrf)*dtn
 
             ! Calculate surface precipitation
             ! Apply sedimenting flux at surface to get rho*dq term
             prec_sed(i) = - q_sed_flux(1)*dtn/dz
             
-            ! This has been moved outside to interface routine
-            ! ! As a final check enforce q must be >= 0.0
-            ! do k = 1,nrf
-            !     q(i,j,k) = max(0.,q(i,j,k))
-            ! end do
+            ! As a final check enforce q must be >= 0.0
+            do k = 1,nrf
+                q(i,k) = max(0.,q(i,k))
+            end do
         end do
         ! End of loop over columns
 
-    deallocate(t)
-    deallocate(q)
-
     end subroutine nn_convection_flux