From 6c09d5254d52ffeab120b6836f0f8c9d0d948537 Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Thu, 24 Nov 2022 17:57:48 +0000 Subject: [PATCH] fixup --- aimodel/src/subcommands/train_mono_predict.py | 1 + 1 file changed, 1 insertion(+) diff --git a/aimodel/src/subcommands/train_mono_predict.py b/aimodel/src/subcommands/train_mono_predict.py index 709bef0..349dd1c 100644 --- a/aimodel/src/subcommands/train_mono_predict.py +++ b/aimodel/src/subcommands/train_mono_predict.py @@ -113,6 +113,7 @@ def do_png(args, ai, dataset, model_params): water_predict = tf.math.argmax(water_predict, axis=-1) # [ width, height, bins ] water_actual = tf.squeeze(water[i_batch]) + # [ width, height ] water_actual = tf.math.argmax(water_actual, axis=-1) segmentation_plot(