test loss function

Arno1235 · Mar 20, 2023 · 5a031ad · 5a031ad
1 parent e578b10
commit 5a031ad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/unet_2d.py b/unet_2d.py
@@ -321,7 +321,7 @@ def get_loss(self, logits, labels):
 
         ### TEST ###
 
-        pred_prob = tf.nn.softmax(logits, axis=-1)
+        pred_prob = tf.nn.softmax(flat_logits, axis=-1)
         pred = tf.one_hot(tf.argmax(pred_prob, -1), self.n_class)
 
         flat_pred = tf.reshape(pred, [-1, self.n_class])