Update app.py

app.py

@@ -10,7 +10,7 @@ from datasets import load_dataset
 from tokenizers import ByteLevelBPETokenizer
 import trl
 
-dataset = load_dataset("nroggendorff/openhermes", split="train")#.select(range(int(4e+5)))
+dataset = load_dataset("nroggendorff/openhermes", split="train").select(range(int(4e+5)))
 
 def get_training_corpus():
     for i in range(0, len(dataset), 1000):
@@ -99,8 +99,7 @@ print(dataset['text'][2])
 args = TrainingArguments(
     output_dir="mayo",
     num_train_epochs=2,
-    gradient_accumulation_steps=16,
-    per_device_train_batch_size=32,
+    per_device_train_batch_size=64,
     learning_rate=1e-5,
     save_steps=100000,
     fp16=True,