Set num_tokens to the correct value

BlinkDL · May 26, 2023 · b227471 · b227471
1 parent c9fbc58
commit b227471
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/RWKV-v4/train.py b/RWKV-v4/train.py
@@ -195,7 +195,7 @@
 
  tconf = TrainerConfig(model_type=model_type, max_epochs=n_epoch, batch_size=batch_size,
  learning_rate=lr_init, lr_decay=True, lr_final=lr_final, betas=betas, eps=eps,
- warmup_tokens=warmup_tokens, final_tokens=n_epoch*len(train_dataset)*ctx_len, num_workers=num_workers, epoch_save_frequency=epoch_save_frequency, epoch_save_path=epoch_save_path)
+ warmup_tokens=warmup_tokens, final_tokens=299892736000, num_workers=num_workers, epoch_save_frequency=epoch_save_frequency, epoch_save_path=epoch_save_path)
  m_cfg = types.SimpleNamespace()
  m_cfg.model_type = model_type
  m_cfg.n_layer = n_layer