diff --git a/keras_rcnn/backend/tensorflow_backend.py b/keras_rcnn/backend/tensorflow_backend.py
index 68d38336..00f49212 100644
--- a/keras_rcnn/backend/tensorflow_backend.py
+++ b/keras_rcnn/backend/tensorflow_backend.py
@@ -1,3 +1,5 @@
+import itertools
+
 import keras.backend
 import numpy
 import tensorflow
diff --git a/keras_rcnn/layers/object_detection/_object_proposal.py b/keras_rcnn/layers/object_detection/_object_proposal.py
index fe4946ec..63e49027 100644
--- a/keras_rcnn/layers/object_detection/_object_proposal.py
+++ b/keras_rcnn/layers/object_detection/_object_proposal.py
@@ -9,7 +9,7 @@ class ObjectProposal(keras.engine.topology.Layer):
     """
 
     def __init__(self, proposals, **kwargs):
-        self.output_dim = (None, proposals, 4)
+        self.output_dim = (None, None, 4)
 
         self.proposals = proposals
 
diff --git a/tests/layers/object_detection/test_anchor.py b/tests/layers/object_detection/test_anchor.py
index fec4e377..b386cb2e 100644
--- a/tests/layers/object_detection/test_anchor.py
+++ b/tests/layers/object_detection/test_anchor.py
@@ -1,6 +1,7 @@
 import keras.backend
 
 import keras_rcnn.backend
+import keras_rcnn.layers.object_detection
 
 
 class TestAnchor:
@@ -18,6 +19,13 @@ def test_call(self, anchor_layer, gt_boxes):
 
         assert n_all_bbox == 1764
 
+    # def test_predict(self):
+    #     shape = (None, 5)
+    #
+    #     x = keras.layers.Input(shape)
+    #
+    #     y = keras_rcnn.layers.object_detection.Anchor((14, 14), (224, 224))(x)
+
     def test_regression(self, anchor_layer, gt_boxes):
         _, y_true, inds_inside, _ = anchor_layer.call(gt_boxes)
 
diff --git a/tests/layers/object_detection/test_object_proposal.py b/tests/layers/object_detection/test_object_proposal.py
index 45778e44..e727a19b 100644
--- a/tests/layers/object_detection/test_object_proposal.py
+++ b/tests/layers/object_detection/test_object_proposal.py
@@ -30,22 +30,16 @@ def test_call(self):
 
         y = keras.layers.MaxPooling2D(strides=(2, 2))(y)
 
-        y = keras.layers.Conv2D(256, **options)(y)
-        y = keras.layers.Conv2D(256, **options)(y)
         y = keras.layers.Conv2D(256, **options)(y)
         y = keras.layers.Conv2D(256, **options)(y)
 
         y = keras.layers.MaxPooling2D(strides=(2, 2))(y)
 
-        y = keras.layers.Conv2D(512, **options)(y)
-        y = keras.layers.Conv2D(512, **options)(y)
         y = keras.layers.Conv2D(512, **options)(y)
         y = keras.layers.Conv2D(512, **options)(y)
 
         y = keras.layers.MaxPooling2D(strides=(2, 2))(y)
 
-        y = keras.layers.Conv2D(512, **options)(y)
-        y = keras.layers.Conv2D(512, **options)(y)
         y = keras.layers.Conv2D(512, **options)(y)
 
         y = keras.layers.Conv2D(512, **options)(y)
@@ -60,11 +54,11 @@ def test_call(self):
 
         model.compile("sgd", "mse")
 
-        image = numpy.random.rand(1, 224, 224, 3)
+        image = numpy.random.rand(1, *shape)
 
         prediction = model.predict(image)
 
         assert prediction.shape == (1, 300, 4)
 
     def test_compute_output_shape(self, object_proposal_layer):
-        assert object_proposal_layer.compute_output_shape((14, 14)) == (None, 300, 4)
+        assert object_proposal_layer.compute_output_shape((14, 14)) == (None, None, 4)