Optimize for memory use.

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}