Add warm up

The current code does not use the warm-up step parameter. Should this change? I would suggest either removing the parameter or using transformers.get_constant_schedule_with_warmup (find more information here