  
  
  ############################################################
-@@ -853,8 +864,9 @@ def rpn_graph(feature_map, anchors_per_location, anchor_stride):
+@@ -853,8 +864,8 @@ def rpn_graph(feature_map, anchors_per_location, anchor_stride):
                    activation='linear', name='rpn_class_raw')(shared)
  
      # Reshape to [batch, anchors, 2]

 2 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/mrcnn/model.py b/mrcnn/model.py
index 62cb2b0..1508f2c 100644
--- a/mrcnn/model.py
+++ b/mrcnn/model.py
@@ -110,17 +110,17 @@ def identity_block(input_tensor, kernel_size, filters, stage, block,

-        lambda t: tf.reshape(t, [tf.shape(t)[0], -1, 2]))(x)
+    x = KL.Permute((2,3,1))(x)
+    rpn_class_logits = KL.Reshape((-1, 2))(x)
 
     # Softmax on last dimension of BG/FG.
     rpn_probs = KL.Activation(
@@ -866,7 +877,8 @@ def rpn_graph(feature_map, anchors_per_location, anchor_stride):
                   activation='linear', name='rpn_bbox_pred')(shared)
 
     # Reshape to [batch, anchors, 4]
-    rpn_bbox = KL.Lambda(lambda t: tf.reshape(t, [tf.shape(t)[0], -1, 4]))(x)
+    x = KL.Permute((2,3,1))(x)
+    rpn_bbox = KL.Reshape((-1, 4))(x)
 
     return [rpn_class_logits, rpn_probs, rpn_bbox]
