diff --git a/tests/keras_contrib/layers/test_convolutional.py b/tests/keras_contrib/layers/test_convolutional.py
index c4acadd..65ff23b 100644
--- a/tests/keras_contrib/layers/test_convolutional.py
+++ b/tests/keras_contrib/layers/test_convolutional.py
@@ -86,6 +86,11 @@ def test_sub_pixel_upscaling():
     for scale_factor in [2, 3, 4]:
         input_data = np.random.random((nb_samples, 4 * (scale_factor ** 2), nb_row, nb_col))
 
+        if K.backend() == 'theano':
+            K.set_image_dim_ordering('th')
+        else:
+            K.set_image_dim_ordering('tf')
+
         if K.image_dim_ordering() == 'tf':
             input_data = input_data.transpose((0, 2, 3, 1))