From aed2348a95e76c28b4345e4e0e81fbfbc34f29f2 Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Thu, 20 Oct 2022 15:42:33 +0100 Subject: [PATCH] train_predict: fixup --- aimodel/src/subcommands/train_predict.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/aimodel/src/subcommands/train_predict.py b/aimodel/src/subcommands/train_predict.py index 037441d..52ea72d 100644 --- a/aimodel/src/subcommands/train_predict.py +++ b/aimodel/src/subcommands/train_predict.py @@ -91,7 +91,8 @@ def do_png(args, ai, dataset, model_code): model_params = json.loads(readfile(args.params)) i = 0 - for rainfall, water in batched_iterator(dataset, tensors_in_item=2, batch_size=model_params["batch_size"]): + gen = batched_iterator(dataset, tensors_in_item=2, batch_size=model_params["batch_size"]) + for rainfall, water in gen: water_predict_batch = ai.embed(rainfall) for water_predict in water_predict_batch: # [ width, height, softmax_probabilities ] → [ batch, width, height ]