From 88363e4d30f30b9bcd66e1f8dd5ae8d814d02c4a Mon Sep 17 00:00:00 2001 From: Huayi Zhang Date: Fri, 19 Jun 2020 12:25:02 +0800 Subject: [PATCH] Bugfix for data_format='channels_first' Bugfix for data_format='channels_first' --- models/official/efficientnet/efficientnet_model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models/official/efficientnet/efficientnet_model.py b/models/official/efficientnet/efficientnet_model.py index 9b969f75f..0290330eb 100644 --- a/models/official/efficientnet/efficientnet_model.py +++ b/models/official/efficientnet/efficientnet_model.py @@ -421,7 +421,8 @@ def call(self, inputs, training=True, survival_prob=None): if self._block_args.id_skip: if all( s == 1 for s in self._block_args.strides - ) and inputs.get_shape().as_list()[-1] == x.get_shape().as_list()[-1]: + ) and (inputs.get_shape().as_list()[self._channel_axis] == + x.get_shape().as_list()[self._channel_axis]): # Apply only if skip connection presents. if survival_prob: x = utils.drop_connect(x, training, survival_prob)