diff --git a/lab5/train/train.py b/lab5/train/train.py
index 270f635..9f24cbb 100644
--- a/lab5/train/train.py
+++ b/lab5/train/train.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 import sys
-import tensorflow as tf
+import tensorflow
 from tensorflow.keras import layers
 
 X_train = pd.read_csv('train.csv')
@@ -17,7 +17,7 @@ Y_test = pd.get_dummies(Y_test)
 Y_valid = X_valid.pop('stabf')
 Y_valid = pd.get_dummies(Y_valid)
 
-model = tf.keras.Sequential([
+model = tensorflow.keras.Sequential([
                           layers.Input(shape=(12,)),
                           layers.Dense(32),
                           layers.Dense(16),
@@ -25,12 +25,12 @@ model = tf.keras.Sequential([
 ])
 
 model.compile(
-    loss=tf.keras.losses.BinaryCrossentropy(),
-    optimizer=tf.keras.optimizers.Adam(lr=float(sys.argv[1])),
-    metrics=[tf.keras.metrics.BinaryAccuracy()])
+    loss=tensorflow.keras.losses.BinaryCrossentropy(),
+    optimizer=tensorflow.keras.optimizers.Adam(lr=float(sys.argv[1])),
+    metrics=[tensorflow.keras.metrics.BinaryAccuracy()])
 
 
-history = model.fit(tf.convert_to_tensor(X_train, np.float32),
+history = model.fit(tensorflow.convert_to_tensor(X_train, np.float32),
                     Y_train, epochs=2, validation_data=(X_valid, Y_valid))
 
 model.save('grid-stability-dense.h5')