diff --git a/lab5/train/train.py b/lab5/train/train.py
index a495381..c3c3c65 100644
--- a/lab5/train/train.py
+++ b/lab5/train/train.py
@@ -17,7 +17,7 @@ Y_valid = X_valid.pop('stabf')
 Y_valid = pd.get_dummies(Y_valid)
 
 model = tensorflow.keras.Sequential([
-                          layers.InputLayer(shape=(12,)),
+                          layers.InputLayer(input_shape=(12,)),
                           layers.Dense(32),
                           layers.Dense(16),
                           layers.Dense(2, activation='softmax')