X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=voxel_flow_model.py;fp=voxel_flow_model.py;h=2cca42054555f85e98998ca7dc67d339e77ad456;hb=511fb18a726351f7cfe3d0b0d4602647496f340f;hp=b3f554d59f935ef327a8a1b86d936461988cb779;hpb=bed00c52fa6fbd9bda76f0f4d2fa2e7c0ff08922;p=voxel-flow

diff --git a/voxel_flow_model.py b/voxel_flow_model.py
index b3f554d..2cca420 100755
--- a/voxel_flow_model.py
+++ b/voxel_flow_model.py
@@ -10,6 +10,8 @@ from utils.geo_layer_utils import vae_gaussian_layer
 from utils.geo_layer_utils import bilinear_interp
 from utils.geo_layer_utils import meshgrid
 
+FLAGS = tf.app.flags.FLAGS
+
 class Voxel_flow_model(object):
   def __init__(self, is_train=True):
     self.is_train = is_train
@@ -69,8 +71,8 @@ class Voxel_flow_model(object):
     mask = tf.expand_dims(net[:, :, :, 2], 3)
 
     grid_x, grid_y = meshgrid(256, 256)
-    grid_x = tf.tile(grid_x, [32, 1, 1]) # batch_size = 32
-    grid_y = tf.tile(grid_y, [32, 1, 1]) # batch_size = 32
+    grid_x = tf.tile(grid_x, [FLAGS.batch_size, 1, 1])
+    grid_y = tf.tile(grid_y, [FLAGS.batch_size, 1, 1])
 
     flow = 0.5 * flow