diff --git a/wgan.py b/wgan.py
index 29e94985a88bed9d53f5c50dfe03833f9453c82e..d3e7a36b80cb1008733e60d2bae43253d525000e 100644
--- a/wgan.py
+++ b/wgan.py
@@ -267,7 +267,7 @@ def run(args):
 
 
     sampler = DistributedSampler(dataset, shuffle=True)    
-    train_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler, num_workers=1, pin_memory=False)
+    train_loader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler, num_workers=1, drop_last=True, pin_memory=False)