Fix a tensorflow test bug. (apache#3165)

Length of input_shape isn't always 4.
wweic · May 13, 2019 · 9b44734 · 9b44734
1 parent 82f5a55
commit 9b44734
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
@@ -185,7 +185,7 @@ def _test_pooling_iteration(input_shape, **kwargs):
 def _test_pooling(input_shape, **kwargs):
     _test_pooling_iteration(input_shape, **kwargs)
 
-    if is_gpu_available():
+    if is_gpu_available() and (len(input_shape) == 4):
         input_shape = [input_shape[ii] for ii in (0, 3, 1, 2)]
         kwargs['data_format'] = 'NCHW'
         _test_pooling_iteration(input_shape, **kwargs)