update max_legn initialization in model

src/main.py

@@ -144,6 +144,7 @@ def __init__(self, d_model, n_heads, num_layers, vocab_size, max_len):
         self.n_heads = n_heads
         self.num_layers = num_layers
         self.vocab_size = vocab_size
+        self.max_len = max_len
         self.embedding = nn.Embedding(vocab_size, d_model)
         self.rope = RoPEPositionalEncoding(d_model, max_len)
         self.transformers = nn.ModuleList(

train.py

@@ -8,6 +8,7 @@
 from torch.nn.utils.rnn import pad_sequence
 import gzip
 from transformers import GPT2Tokenizer
+from datasets import load_dataset
 from importlib import reload
 import src.main
 from accelerate import Accelerator