diff --git a/localizer/predict.py b/localizer/predict.py
index b81dd0b..d1fe738 100644
--- a/localizer/predict.py
+++ b/localizer/predict.py
@@ -183,7 +183,7 @@ def show_kernel(k, name, factor):
                                  strides=1,
                                  padding='SAME', name='confidence_map')
 
-        confidence_loc_max = (tf.nn.max_pool(confidence_map, 3, 1, 'SAME', 'NCDHW') - confidence_map) == 0
+        confidence_loc_max = (tf.nn.max_pool(confidence_map, 3, 1, 'SAME', 'NDHWC') - confidence_map) == 0
         confidence_thr = tf.math.logical_and(confidence_loc_max, confidence_map >= self._cfg['confidence_thr'])
         confidence = confidence_map * tf.cast(confidence_thr, tf.float32)