diff --git a/generate.py b/generate.py
index 3f289de..de3fabc 100644
--- a/generate.py
+++ b/generate.py
@@ -43,7 +43,7 @@
             embedding_dim=256,
             vocab_size=par.vocab_size,
             num_layer=6,
-            max_seq=2048,
+            max_seq=max_seq,
             dropout=0.2,
             debug=False, loader_path=load_path)
 else:
diff --git a/model.py b/model.py
index 26e3c01..a74a8bd 100644
--- a/model.py
+++ b/model.py
@@ -39,9 +39,9 @@ def __init__(self, embedding_dim=256, vocab_size=388+2, num_layer=6,
             self.load_ckpt_file(loader_path)
 
     def call(self, inputs, targets, training=None, eval=None, src_mask=None, trg_mask=None, lookup_mask=None):
-        encoder, weight_encoder = self.Encoder(inputs, training=training, mask=src_mask)
+        encoder, weight_encoder = self.Encoder(inputs, training=training, mask=src_mask, max_len=self.max_seq)
         decoder, weights = self.Decoder(
-            targets, enc_output=encoder, training=training, lookup_mask=lookup_mask, mask=trg_mask
+            targets, enc_output=encoder, training=training, lookup_mask=lookup_mask, mask=trg_mask, max_len=self.max_seq
         )
 
         fc = self.fc(decoder)
@@ -275,7 +275,7 @@ def __init__(self, embedding_dim=256, vocab_size=388+2, num_layer=6,
 
         self.Decoder = Encoder(
             num_layers=self.num_layer, d_model=self.embedding_dim,
-            input_vocab_size=self.vocab_size, rate=dropout, max_len=max_seq)
+            input_vocab_size=self.vocab_size, rate=dropout, max_len=self.max_seq)
         self.fc = keras.layers.Dense(self.vocab_size, activation=None, name='output')
 
         self._set_metrics()
diff --git a/train.py b/train.py
index 767586b..9c781fd 100644
--- a/train.py
+++ b/train.py
@@ -41,7 +41,8 @@
 
 
 # load data
-dataset = Data('dataset/processed')
+#dataset = Data('dataset/processed')
+dataset = Data(pickle_dir)
 print(dataset)