diff --git a/mlspm/image/models.py b/mlspm/image/models.py
index 27c0d47..b62aa38 100644
--- a/mlspm/image/models.py
+++ b/mlspm/image/models.py
@@ -52,7 +52,7 @@ class AttentionUNet(nn.Module):
 
     def __init__(
         self,
-        z_in: int = 10,
+        z_in: int,
         n_in: int = 1,
         n_out: int = 3,
         in_channels: int = 1,