From 3d271e9c44e5831bbd96217d87140673b5ca997b Mon Sep 17 00:00:00 2001 From: Yin Li Date: Tue, 18 May 2021 14:16:40 -0400 Subject: [PATCH] Add Eulerian shape logging --- map2map/train.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/map2map/train.py b/map2map/train.py index 2388f7a..83ffddc 100644 --- a/map2map/train.py +++ b/map2map/train.py @@ -266,10 +266,12 @@ def train(epoch, loader, model, criterion, input = resample(input, model.module.scale_factor, narrow=False) input, output, target = narrow_cast(input, output, target) if batch <= 5 and rank == 0: - print('narrowed shape :', output.shape, flush=True) + print('narrowed shape :', output.shape) lag_out, lag_tgt = output, target eul_out, eul_tgt = lag2eul([lag_out, lag_tgt], **args.misc_kwargs) + if batch <= 5 and rank == 0: + print('Eulerian shape :', eul_out.shape, flush=True) lag_loss = criterion(lag_out, lag_tgt) eul_loss = criterion(eul_out, eul_tgt)