code: def record_lr(
      model: tf.keras.Model,
      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"
          )
  
      # Update param groups & LR
      gamma = (end_lr / start_lr) ** (1 / (num_it - 1))
      optimizer.learning_rate = start_lr
  
      lr_recorder = [start_lr * gamma**idx for idx in range(num_it)]
      loss_recorder = []
  
      for batch_idx, (images, targets) in enumerate(train_loader):
          images = batch_transforms(images)
  
          # Forward, Backward & update
          with tf.GradientTape() as tape:
              out = model(images, training=True)
              train_loss = (
                  tf.nn.sparse_softmax_cross_entropy_with_logits(
                      targets, out
                  )
              )
          grads = tape.gradient(train_loss, model.trainable_weights)
  
          if amp:
              grads = optimizer.get_unscaled_gradients(grads)
          optimizer.apply_gradients(zip(grads, model.trainable_weights))
  
          optimizer.learning_rate = optimizer.learning_rate * gamma
  
          # Record
          train_loss = train_loss.numpy()
          if np.any(np.isnan(train_loss)):
              if batch_idx == 0:
                  raise ValueError("loss value is NaN or inf.")
              else:
                  break
          loss_recorder.append(train_loss.mean())
          # Stop after the number of iterations
          if batch_idx + 1 == num_it:
              break
  
      return lr_recorder[: len(loss_recorder)], loss_recorder
  
file_path: doctr/references/classification/train_tensorflow.py
imports: import os
  import datetime
  import multiprocessing
  import time
  import numpy
  import tensorflow
  from tensorflow.keras import mixed_precision
  from tqdm.auto import tqdm
  from doctr.models import login_to_hub, push_to_hf_hub
  from doctr import transforms
  from doctr.datasets import VOCABS, CharacterGenerator, DataLoader
  from doctr.models import classification
  from doctr.models.utils import export_model_to_onnx
  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