Fix shape issue

Ref: melodyguan#29

src/cifar10/general_child.py

@@ -291,13 +291,13 @@ def _enas_layer(self, layer_id, prev_layers, start_idx, out_filters, is_training
           y = self._pool_branch(inputs, is_training, out_filters, "max",
                                 start_idx=0)
         branches[tf.equal(count, 5)] = lambda: y
-      out = tf.case(branches, default=lambda: tf.constant(0, tf.float32),
-                    exclusive=True)
 
       if self.data_format == "NHWC":
-        out.set_shape([None, inp_h, inp_w, out_filters])
+        out_shape = [self.batch_size, inp_h, inp_w, out_filters]
       elif self.data_format == "NCHW":
-        out.set_shape([None, out_filters, inp_h, inp_w])
+        out_shape = [self.batch_size, out_filters, inp_h, inp_w]
+      out = tf.case(branches, default=lambda: tf.constant(0, tf.float32, shape=out_shape),
+                    exclusive=True)
     else:
       count = self.sample_arc[start_idx:start_idx + 2 * self.num_branches]
       branches = []