Fixed wrong shapes with BS=1

hXl3s · hXl3s · commit 43d061ae8727 · 2021-03-11T10:31:06.000+01:00
diff --git a/TensorFlow/Classification/ConvNets/model/resnet.py b/TensorFlow/Classification/ConvNets/model/resnet.py
@@ -187,12 +187,9 @@ def __call__(self, features, labels, mode, params):
                 reuse=False,
                 use_final_conv=params['use_final_conv']
             )
-            
-            if mode!=tf.estimator.ModeKeys.PREDICT:
-                logits = tf.squeeze(logits)
 
-            if mode!=tf.estimator.ModeKeys.PREDICT:
-                logits = tf.squeeze(logits)
+            if params['use_final_conv']:
+                logits = tf.squeeze(logits, axis=[-2, -1])
 
             y_preds = tf.argmax(logits, axis=1, output_type=tf.int32)
 
