code: def record_lr(
      model: torch.nn.Module,
      train_loader: DataLoader,
      batch_transforms,
      optimizer,
      start_lr: float = 1e-7,
      end_lr: float = 1,
      num_it: int = 100,
      amp: bool = False,
  ):
      """Gridsearch the optimal learning rate for the training.
      Adapted from https://github.com/frgfm/Holocron/blob/master/holocron/trainer/core.py
      """
      if num_it > len(train_loader):
          raise ValueError(
              "the value of `num_it` needs to be lower than the number of available batches"
          )
  
      model = model.train()
      # Update param groups & LR
      optimizer.defaults["lr"] = start_lr
      for pgroup in optimizer.param_groups:
          pgroup["lr"] = start_lr
  
      gamma = (end_lr / start_lr) ** (1 / (num_it - 1))
      scheduler = MultiplicativeLR(optimizer, lambda step: gamma)
  
      lr_recorder = [start_lr * gamma**idx for idx in range(num_it)]
      loss_recorder = []
  
      if amp:
          scaler = torch.cuda.amp.GradScaler()
  
      for batch_idx, (images, targets) in enumerate(train_loader):
          if torch.cuda.is_available():
              images = images.cuda()
  
          images = batch_transforms(images)
  
          # Forward, Backward & update
          optimizer.zero_grad()
          if amp:
              with torch.cuda.amp.autocast():
                  train_loss = model(images, targets)["loss"]
              scaler.scale(train_loss).backward()
              # Gradient clipping
              scaler.unscale_(optimizer)
              torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
              # Update the params
              scaler.step(optimizer)
              scaler.update()
          else:
              train_loss = model(images, targets)["loss"]
              train_loss.backward()
              torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
              optimizer.step()
          # Update LR
          scheduler.step()
  
          # Record
          if not torch.isfinite(train_loss):
              if batch_idx == 0:
                  raise ValueError("loss value is NaN or inf.")
              else:
                  break
          loss_recorder.append(train_loss.item())
          # Stop after the number of iterations
          if batch_idx + 1 == num_it:
              break
  
      return lr_recorder[: len(loss_recorder)], loss_recorder
  
imports: import os
  import datetime
  import hashlib
  import logging
  import multiprocessing
  import time
  from pathlib import Path
  import numpy
  import torch
  import wandb
  from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR, OneCycleLR
  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
  from torchvision.transforms.v2 import Compose, GaussianBlur, Normalize, RandomGrayscale, RandomPerspective, RandomPhotometricDistort
  from tqdm.auto import tqdm
  from doctr import transforms
  from doctr.datasets import VOCABS, RecognitionDataset, WordGenerator
  from doctr.models import login_to_hub, push_to_hf_hub, recognition
  from doctr.utils.metrics import TextMatch
  from utils import EarlyStopper, plot_recorder, plot_samples
label_desc: Use image node_id 1 for calling a global function with example usage: record_lr(model, train_loader, batch_transforms, optimizer, start_lr, end_lr, num_it, amp) and returns: lr_recorder,  loss_recorder