code: def _find_lr(self, train_data):
      lr = 1e-4  # good starting point as search escalates
      lrs = deque([5e-4, 1e-3, 2e-3, 3e-3, 5e-3, 1e-2, 5e-2, 1e-1])
      starting_model = deepcopy(self.model)
      criterion = self._select_criterion()
      scaler = GradScaler()
  
      running_losses = deque(maxlen=self.loss_hist_len)
      lr_log = deque(maxlen=self.loss_hist_len)
      best_model = self.model
      stop = False
  
      n_steps = 10
      cum_loss = 0
  
      while stop is False:
          # overfit learning on first n_steps samples (biased, but we only want an intuition on what LR is decent)
          dl = DataLoader(
              train_data,
              batch_size=min(
                  len(train_data.data_frame), 32, self.batch_size
              ),
              shuffle=False,
          )
          dl_iter = iter(dl)
          self.model = deepcopy(starting_model)
          self.model.train()
          optimizer = self._select_optimizer(self.model, lr=lr)
  
          for i in range(n_steps):
              try:
                  X, Y = next(dl_iter)
              except StopIteration:
                  dl_iter = iter(dl)
                  X, Y = next(dl_iter)
  
              X = X.to(self.model.device)
              Y = Y.to(self.model.device)
  
              with LightwoodAutocast():
                  optimizer.zero_grad()
                  Yh = self._net_call(X)
                  loss = criterion(Yh, Y)
                  if LightwoodAutocast.active:
                      scaler.scale(loss).backward()
                      scaler.step(optimizer)
                      scaler.update()
                  else:
                      loss.backward()
                      optimizer.step()
              cum_loss += loss.item()
  
          log.info(f"Loss of {cum_loss} with learning rate {lr}")
          running_losses.append(cum_loss)
          lr_log.append(lr)
          cum_loss = 0
          lr = lrs.popleft()
          if len(lrs) == 0:
              stop = True
  
          # store model if best so far
          inv_running_losses = list(running_losses)[
              ::-1
          ]  # invert so when tied we pick the most aggresive LR
          best_loss_idx = np.nanargmin(
              inv_running_losses
          )  # nanargmin ignores nans that may arise
          if best_loss_idx == 0:
              best_model = deepcopy(
                  self.model
              )  # store model for slight time savings
              best_loss_lr = lr_log[-1]
  
      lr = best_loss_lr
      log.info(f"Found learning rate of: {lr}")
      return lr, best_model
  
imports: import time
  from copy import deepcopy
  from collections import deque
  from typing import Dict, List, Optional
  import torch
  import numpy
  import pandas
  from torch import nn
  import torch_optimizer
  from sklearn.metrics import r2_score
  from torch.cuda.amp import GradScaler
  from torch.utils.data import DataLoader
  from torch.nn.modules.loss import MSELoss
  from torch.optim.optimizer import Optimizer
  from type_infer.dtype import dtype
  from lightwood.helpers.log import log
  from lightwood.encoder.base import BaseEncoder
  from lightwood.helpers.torch import LightwoodAutocast
  from lightwood.data.encoded_ds import EncodedDs
  from lightwood.mixer.base import BaseMixer
  from lightwood.mixer.helpers.ar_net import ArNet
  from lightwood.mixer.helpers.default_net import DefaultNet
  from lightwood.api.types import PredictionArguments
  from lightwood.mixer.helpers.transform_corss_entropy_loss import TransformCrossEntropyLoss
label_desc: Use image node_id 5 for calling the Neural obj's underlying member method code with example usage: obj._find_lr(train_data) and returns: lr,  best_model