Update train.py

train.py

@@ -124,13 +124,13 @@ config = {"train_batch_size": 2,
           "shuffle_buffer": 1000,
           "learning_rate": 5e-4,
           "lr_scheduler_type": "cosine",
-          "num_warmup_steps": 0,
+          "num_warmup_steps": 2000,
           "gradient_accumulation_steps": 1,
-          "max_train_steps": 15,
-          "max_eval_steps": 15,
+          "max_train_steps": 150000,
+          "max_eval_steps": -1,
           "seq_length": 1024,
           "seed": 1,
-          "save_checkpoint_steps": 10}
+          "save_checkpoint_steps": 15000}
 args = Namespace(**config, **acc_state)
 samples_per_step = accelerator.state.num_processes * args.train_batch_size
 set_seed(args.seed)