diff --git a/aimodel/src/subcommands/train_predict.py b/aimodel/src/subcommands/train_predict.py
index 0d2b8ba..fb94adb 100644
--- a/aimodel/src/subcommands/train_predict.py
+++ b/aimodel/src/subcommands/train_predict.py
@@ -99,7 +99,7 @@ def do_png(args, ai, dataset, model_code):
 		rainfall, water = item
 		
 		water_predict_batch = ai.embed(rainfall)
-		water = tf.unstack(tf.squeeze(water), axis=0)
+		water = tf.squeeze(tf.unstack(water, axis=0))
 		
 		i_batch = 0
 		for water_predict in water_predict_batch: