diff --git a/IUM_05-predict.py b/IUM_05-predict.py
index 60681b2..1d5b603 100644
--- a/IUM_05-predict.py
+++ b/IUM_05-predict.py
@@ -2,18 +2,30 @@ import pandas as pd
 import numpy as np
 import tensorflow as tf
 
+# Load the test data
 test_data = pd.read_csv('./beer_reviews_test.csv')
 X_test = test_data[['review_aroma', 'review_appearance', 'review_palate', 'review_taste']]
 y_test = test_data['review_overall']
 
+# Load the model
 model = tf.keras.models.load_model('beer_review_sentiment_model.h5')
 
+# Preprocess the test data
 tokenizer = tf.keras.preprocessing.text.Tokenizer(num_words=10000)
 tokenizer.fit_on_texts(X_test)
 X_test_seq = tokenizer.texts_to_sequences(X_test)
 X_test_pad = tf.keras.preprocessing.sequence.pad_sequences(X_test_seq, maxlen=100)
 
+# Make predictions
 predictions = model.predict(X_test_pad)
 
-results = pd.DataFrame({'Predictions': predictions.flatten(), 'Actual': y_test})
+# Check the shape of the predictions
+print(f'Predictions shape: {predictions.shape}')
+
+# If predictions have more than one dimension, select only one dimension
+if len(predictions.shape) > 1:
+    predictions = predictions[:, 0]
+
+# Save predictions and actual data to a CSV file
+results = pd.DataFrame({'Predictions': predictions, 'Actual': y_test})
 results.to_csv('beer_review_sentiment_predictions.csv', index=False)
\ No newline at end of file