Udpated concat axis to match image_data_format in keras (#2946)

microsoft · Oct 19, 2020 · 3ffd105 · 3ffd105
1 parent 58873c4
commit 3ffd105
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/sdk/pynni/nni/nas/tensorflow/mutator.py b/src/sdk/pynni/nni/nas/tensorflow/mutator.py
@@ -66,7 +66,12 @@ def _tensor_reduction(self, reduction_type, tensor_list):
         if reduction_type == 'mean':
             return sum(tensor_list) / len(tensor_list)
         if reduction_type == 'concat':
-            return tf.concat(tensor_list, axis=0)
+            image_data_format = tf.keras.backend.image_data_format()
+            if image_data_format == "channels_first":
+                axis = 0
+            else:
+                axis = -1
+            return tf.concat(tensor_list, axis=axis)
         raise ValueError('Unrecognized reduction policy: "{}'.format(reduction_type))
 
     def _get_decision(self, mutable):