diff --git a/train.py b/train.py
index 767586b..29e08c4 100644
--- a/train.py
+++ b/train.py
@@ -41,7 +41,7 @@
 
 
 # load data
-dataset = Data('dataset/processed')
+dataset = Data(pickle_dir)
 print(dataset)