From bbc21f48993545d49e9d0626c965937c66ea8f68 Mon Sep 17 00:00:00 2001 From: Jan Kantert Date: Thu, 22 Dec 2022 16:09:42 +0100 Subject: [PATCH] fix inference on CPU Previously you would get "Default Pooling3DOp only supports NDHWC on device type CPU". --- localizer/predict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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)