diff --git a/create-dataset.py b/create-dataset.py
index 209d2ac..97fef25 100644
--- a/create-dataset.py
+++ b/create-dataset.py
@@ -10,7 +10,6 @@ from sklearn.model_selection import train_test_split
 
 
 def download_kaggle_dataset():
-    print(sys.argv)
     os.environ["KAGGLE_USERNAME"] = "vskyper"
     os.environ["KAGGLE_KEY"] = sys.argv[1]
     kaggle = KaggleApi()
@@ -118,7 +117,7 @@ def main():
     )
 
     X_train, X_test, y_train, y_test = split_whole_data(df)
-    save_whole_data(X_train, X_test, y_train, y_test)
+    save_whole_data(df, X_train, X_test, y_train, y_test)
 
 
 if __name__ == "__main__":