code: def _max_fit(
      self,
      train_dl,
      dev_dl,
      criterion,
      optimizer,
      scaler,
      stop_after,
      return_model_after,
  ):
      epochs_to_best = 0
      best_dev_error = pow(2, 32)
      running_errors = deque(maxlen=self.loss_hist_len)
      best_model = self.model
  
      for epoch in range(1, return_model_after + 1):
          self.model = self.model.train()
          running_losses: List[float] = []
          for i, (X, Y) in enumerate(train_dl):
              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()
  
              running_losses.append(loss.item())
              if (time.time() - self.started) > stop_after:
                  break
  
          train_error = np.mean(running_losses)
          epoch_error = self._error(dev_dl, criterion)
          running_errors.append(epoch_error)
          log.info(f"Loss @ epoch {epoch}: {epoch_error}")
  
          if (
              np.isnan(train_error)
              or np.isnan(running_errors[-1])
              or np.isinf(train_error)
              or np.isinf(running_errors[-1])
          ):
              break
  
          if best_dev_error > running_errors[-1]:
              best_dev_error = running_errors[-1]
              best_model = deepcopy(self.model)
              epochs_to_best = epoch
  
          # manually set epoch limit
          if self.n_epochs is not None:
              if epoch > self.n_epochs:
                  break
  
          # automated early stopping
          else:
              if len(running_errors) >= self.loss_hist_len:
                  delta_mean = np.average(
                      [
                          running_errors[-i - 1] - running_errors[-i]
                          for i in range(len(running_errors) - 1)
                      ],
                      weights=[
                          (1 / 2) ** i
                          for i in range(len(running_errors) - 1)
                      ],
                  )
                  if delta_mean >= 0:
                      break
              elif (time.time() - self.started) > stop_after:
                  break
              elif running_errors[-1] < 0.0001 or train_error < 0.0001:
                  break
  
      if np.isnan(best_dev_error):
          best_dev_error = pow(2, 32)
      return best_model, epochs_to_best, best_dev_error
  
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 6 for calling the Neural obj's underlying member method code with example usage: obj._max_fit(train_dl, dev_dl, criterion, optimizer, scaler, stop_after, return_model_after) and returns: best_model,  epochs_to_best,  best_dev_error