diff --git a/wgan/wgan.py b/wgan/wgan.py
index 4a8bb86..402f77c 100644
--- a/wgan/wgan.py
+++ b/wgan/wgan.py
@@ -39,7 +39,7 @@ class WGAN():
         self.generator = self.build_generator()
 
         # The generator takes noise as input and generated imgs
-        z = Input(shape=(100,))
+        z = Input(shape=(self.latent_dim,))
         img = self.generator(z)
 
         # For the combined model we will only train the generator