support batch inpainting with shared mask in each batch

WxpHB · Mar 11, 2018 · bd73aeb · bd73aeb
1 parent 424af23
commit bd73aeb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/inpaint_model.py b/inpaint_model.py
@@ -297,7 +297,7 @@ def build_server_graph(self, batch_data, reuse=False, is_training=False):
         """
         # generate mask, 1 represents masked point
         batch_raw, masks_raw = tf.split(batch_data, 2, axis=2)
-        masks = tf.cast(masks_raw[:, :, :, 0:1] > 127.5, tf.float32)
+        masks = tf.cast(masks_raw[0:1, :, :, 0:1] > 127.5, tf.float32)
 
         batch_pos = batch_raw / 127.5 - 1.
         batch_incomplete = batch_pos * (1. - masks)