Update train.py

tensorlayer · Jul 11, 2019 · 091e11c · 091e11c
1 parent b90cb20
commit 091e11c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/train.py b/train.py
@@ -23,6 +23,8 @@ def train():
 
     for epoch in range(flags.n_epoch):
         for step, batch_images in enumerate(images):
+            if batch_images.shape[0] != flags.batch_size: # if the remaining data in this epoch < batch_size
+                break
             step_time = time.time()
             with tf.GradientTape(persistent=True) as tape:
                 # z = tf.distributions.Normal(0., 1.).sample([flags.batch_size, flags.z_dim]) #tf.placeholder(tf.float32, [None, z_dim], name='z_noise')