dataset keep in memory

scripts/train.py

@@ -33,8 +33,8 @@ def main():
     tokenizer.pad_token = tokenizer.eos_token
     
     # Load dataset
-    dataset = load_dataset("allenai/tulu-v2-sft-mixture", split="train")
-    
+    dataset = load_dataset("allenai/tulu-v2-sft-mixture", split="train", keep_in_memory=True)
+
     def tokenize_function(examples):
         return tokenizer(
             examples["text"],