OMG I found it through adding some prints to debug. The problem is in my yolo_filter_boxes implementation. When I specified the last dimension in tf.math.argmax() I used axis=3, which worked with the shape of inputs at that time. But in section 3.5 yolo_outputs has another dimension (totally 5 dimensions). I fixed my code with axis=-1 and all is well 
1 Like