diff --git a/tests/keras_contrib/layers/test_convolutional.py b/tests/keras_contrib/layers/test_convolutional.py
index 483d5cb..1d4bd68 100644
--- a/tests/keras_contrib/layers/test_convolutional.py
+++ b/tests/keras_contrib/layers/test_convolutional.py
@@ -121,7 +121,7 @@ def test_cosineconvolution_2d():
     elif dim_ordering == 'tf':
         X = np.random.randn(1, 5, 5, 3)
         input_dim = (5, 5, 3)
-        W0 = X[:, ::-1, ::-1, :]
+        W0 = X[0, ::-1, ::-1, :, None]
 
     model = Sequential()
     model.add(convolutional.CosineConvolution2D(1, 5, 5, bias=True, input_shape=input_dim))