From e6f5706175129afb11a7df09a29aa486a3974e4b Mon Sep 17 00:00:00 2001
From: "Daniel.Frisinghelli" <daniel.frisinghelli@eurac.edu>
Date: Mon, 22 Nov 2021 12:09:07 +0100
Subject: [PATCH] Optimize for memory use.
---
climax/core/predict.py | 8 +++-----
1 file changed, 3 insertions(+), 5 deletions(-)
diff --git a/climax/core/predict.py b/climax/core/predict.py
index bbefa41..bd29e2f 100644
--- a/climax/core/predict.py
+++ b/climax/core/predict.py
@@ -74,12 +74,10 @@ def predict_ERA5(net, ERA5_ds, predictand, loss, batch_size=16, **kwargs):
prob = torch.sigmoid(torch.as_tensor(target[:, 0, ...].squeeze(),
dtype=torch.float32)).numpy()
- # shape and scale parameters
- shape = target[:, 1, ...].squeeze()
- scale = target[:, 2, ...].squeeze()
-
# precipitation amount
- pr = loss.predict(prob, shape, scale)
+ pr = loss.predict(prob, target[:, 1, ...].squeeze(),
+ target[:, 2, ...].squeeze())
+ del target # clear memory
# precipitation probability and amount
ds = {'prob': prob, 'precipitation': pr}
--
GitLab