code: def run_test(
      scheduler,
      num_epochs: int,
      num_batches: int,
      return_all_lrs: Optional[bool] = False,
  ) -> Union[List, float]:
      end_lr = [] if return_all_lrs else 0.0
      curr_iter = 0
      for ep in range(num_epochs):
          for _ in range(num_batches):
              lr = scheduler.get_lr(ep, curr_iter=curr_iter)
              curr_iter += 1
  
          # keep only epoch-wise LRs
          if return_all_lrs:
              end_lr.append(lr)
          else:
              end_lr = lr
  
      return end_lr
  
file_path: ml-cvnets/tests/test_scheduler.py
imports: import random
  import sys
  from typing import List, Optional, Union
  import numpy
  import pytest
  from optim.scheduler import build_scheduler
  from tests.configs import get_config
label_desc: Use image node_id 1 for calling a global function with example usage: run_test(scheduler, num_epochs, num_batches, return_all_lrs) and returns: end_lr