Other answers have already covered how to implement your request exactly. But I wanted to mention that Pytorch already has builtin LR warmup options & they can likely achieve comparable (or better!) results with less code!

Builtin Options:

• RAdam, torch.optim.RAdam() : Somewhat shorter (implicit) warmup period specialized for Adam. It can achieve optimal results of Adam with less LR tuning.
• OneCycleLR, torch.optim.lr_scheduler.OneCycleLR() : Applies warmup & decay. Good for high performance training (e.g. scaling LR & batch size to many GPUs).