From 4d23f120273f305d07b4d5a7f12f197eeef8dd7f Mon Sep 17 00:00:00 2001 From: jatkinson1000 Date: Mon, 24 Jun 2024 03:55:54 -0600 Subject: [PATCH] Move reading of NN to initialisation routine for gw_drag. --- src/physics/cam/gw_drag.F90 | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/physics/cam/gw_drag.F90 b/src/physics/cam/gw_drag.F90 index 85819302e9..bfe4cacfff 100644 --- a/src/physics/cam/gw_drag.F90 +++ b/src/physics/cam/gw_drag.F90 @@ -198,6 +198,7 @@ module gw_drag logical :: gw_convect_dp_ml = .false. logical :: gw_convect_dp_ml_compare = .false. character(len=132) :: gw_convect_dp_ml_net + type(torch_module) :: gw_convect_dp_nn !========================================================================== contains @@ -980,6 +981,12 @@ subroutine gw_init() end if + ! Set up neccessary attributes if using ML scheme for convective drag + if ((gw_convect_dp_ml == 'on') .or. (gw_convect_dp_ml == 'bothon')) then + ! Load the convective drag net from TorchScript file + gw_convect_dp_nn = torch_module_load(gw_convect_dp_ml_net) + endif + if (use_gw_convect_sh) then ttend_sh_idx = pbuf_get_index('TTEND_SH') @@ -1417,9 +1424,6 @@ subroutine gw_tend(state, pbuf, dt, ptend, cam_in, flx_heat) !------------------------------------------------------------------------ type(torch_tensor), dimension(1) :: in_tensors - type(torch_module) :: gw_convect_dp_nn - - gw_convect_dp_nn = torch_module_load(gw_convect_dp_ml_net) if (masterproc) then write(iulog,*) 'Made a torch tensor and loaded a net!'