code: def warmup_lr_schedule(optimizer, step, max_step, init_lr, max_lr):
      """Warmup the learning rate"""
      lr = min(max_lr, init_lr + (max_lr - init_lr) * step / max_step)
      for param_group in optimizer.param_groups:
          param_group["lr"] = lr
  
imports: import math
  import numpy
  import io
  import os
  import time
  from collections import defaultdict, deque
  import datetime
  import torch
  import torch.distributed
label_desc: Use image node_id 2 for calling a global function with example usage: warmup_lr_schedule(optimizer, step, max_step, init_lr, max_lr) without return types