diff --git a/models/cluster_gan/clustergan.py b/models/cluster_gan/clustergan.py
index 213e4e5..38ecfbd 100644
--- a/models/cluster_gan/clustergan.py
+++ b/models/cluster_gan/clustergan.py
@@ -280,7 +280,7 @@ def execute(self, img):
         #  Train Generator
         # -----------------
 
-        if ((i % n_skip_iter) == 0):
+        if ((i % n_skip_iter) == 0) or real_imgs.shape[0] != batch_size:
             (enc_gen_zn, enc_gen_zc, enc_gen_zc_logits) = encoder(gen_imgs)
             zn_loss = mse_loss(enc_gen_zn, zn)
             zc_loss = xe_loss(enc_gen_zc_logits, zc_idx)