From 48d29de6084b270194d29e717b9753a58beb43cd Mon Sep 17 00:00:00 2001 From: mwang Date: Thu, 18 Mar 2021 20:11:08 +0800 Subject: [PATCH] add data sync of GOMO --- model_zoo/research/hpc/ocean_model/train.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/model_zoo/research/hpc/ocean_model/train.py b/model_zoo/research/hpc/ocean_model/train.py index 4f1bb6a37bb..be5e4544edc 100644 --- a/model_zoo/research/hpc/ocean_model/train.py +++ b/model_zoo/research/hpc/ocean_model/train.py @@ -46,8 +46,15 @@ if __name__ == "__main__": # define grid and init variable update net_init = GOMO_init(im, jm, kb, stencil_width) + init_res = net_init(dx, dy, dz, uab, vab, elb, etb, sb, tb, ub, vb, h, w, vfluxf, zz, fsm) + for res_tensor in init_res: + if isinstance(res_tensor, (list, tuple)): + for rt in res_tensor: + rt.data_sync(True) + else: + res_tensor.data_sync(True) ua, va, el, et, etf, d, dt, l, q2b, q2lb, kh, km, kq, aam, w, q2, q2l, t, s, u, v, cbc, rmean, rho, x_d, y_d, z_d\ - = net_init(dx, dy, dz, uab, vab, elb, etb, sb, tb, ub, vb, h, w, vfluxf, zz, fsm) + = init_res # define GOMO model Model = GOMO(im=im, jm=jm, kb=kb, stencil_width=stencil_width, variable=variable, x_d=x_d, y_d=y_d, z_d=z_d,