revert the scheduler change
@@ -14,7 +14,7 @@ import torch
from tqdm import tqdm, trange
from itertools import cycle
-from transformers import AdamW, get_linear_schedule_with_warmup
+from transformers import AdamW, WarmupLinearSchedule
from transformers.modeling_bert import BERT_PRETRAINED_MODEL_ARCHIVE_MAP
from transformers.modeling_distilbert import DISTILBERT_PRETRAINED_MODEL_ARCHIVE_MAP
from transformers.modeling_roberta import ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
@@ -150,9 +150,7 @@ class Transformer:
optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=adam_epsilon)
if t_total != -1 and scheduler is None:
- scheduler = get_linear_schedule_with_warmup(
- optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
- )
+ scheduler = WarmupLinearSchedule(optimizer, warmup_steps=warmup_steps, t_total=t_total)
if fp16:
try:
