code: def fit(
      self,
      train_data_loader,
      test_data_loader,
      desired_error,
      max_time,
      callback=None,
      eval_every_x_epochs=1,
      max_unimproving_models=10,
      custom_train_func=None,
      custom_test_func=None,
  ):
      started = time.time()
      epoch = 0
      lowest_test_error = None
      last_test_error = None
      test_error_delta_buff = []
  
      keep_training = True
  
      while keep_training:
          epoch += 1
          running_loss = 0.0
          error = 0
          self.model = self.model.train()
          for i, data in enumerate(train_data_loader, 0):
              if custom_train_func is None:
                  input, real = data
  
                  with LightwoodAutocast():
                      if self.input_encoder is not None:
                          input = self.input_encoder(input)
                          if len(input.shape) < 2:
                              input = input.unsqueeze(-1)
                      if self.output_encoder is not None:
                          real = self.output_encoder(real)
  
                      input = input.to(self.device)
                      real = real.to(self.device)
  
                      predicted = self.model(input)
                      loss = self.loss_criterion(predicted, real)
                  loss.backward()
                  self.optimizer.step()
  
                  if self.scheduler is not None:
                      self.scheduler.step()
  
                  self.optimizer.zero_grad()
              else:
                  loss = custom_train_func(self.model, data, self)
  
              running_loss += loss.item()
              error = running_loss / (i + 1)
  
          # end of epoch checks
          log.debug(
              f"Categorical AutoEncoder train loss at epoch {epoch}: {round(error, 9)}"
          )
          if epoch % eval_every_x_epochs == 0:
              if test_data_loader is not None:
                  test_running_loss = 0.0
                  test_error = 0
                  self.model = self.model.eval()
                  real_buff = []
                  predicted_buff = []
  
                  for i, data in enumerate(test_data_loader, 0):
                      if custom_test_func is None:
                          input, real = data
  
                          if self.input_encoder is not None:
                              input = self.input_encoder(input)
                          if self.output_encoder is not None:
                              real = self.output_encoder(real)
  
                          input = input.to(self.device)
                          real = real.to(self.device)
  
                          with torch.no_grad():
                              predicted = self.model(input)
  
                          real_buff.append(real.tolist())
                          predicted_buff.append(predicted.tolist())
  
                          loss = self.loss_criterion(predicted, real)
                      else:
                          with torch.no_grad():
                              loss = custom_test_func(
                                  self.model, data, self
                              )
  
                      log.debug(
                          f"Categorical AutoEncoder val loss at epoch {epoch}: {round(loss, 9)}"
                      )
                      test_running_loss += loss.item()
                      test_error = test_running_loss / (i + 1)
              else:
                  test_error = error
                  real_buff = None
                  predicted_buff = None
  
              if (
                  lowest_test_error is None
                  or test_error < lowest_test_error
              ):
                  lowest_test_error = test_error
                  self.best_model = copy.deepcopy(self.model).to("cpu")
  
              if last_test_error is None:
                  test_error_delta_buff.append(0)
              else:
                  test_error_delta_buff.append(
                      last_test_error - test_error
                  )
  
              last_test_error = test_error
  
              if (time.time() - started) > max_time:
                  keep_training = False
  
              if lowest_test_error < desired_error:
                  keep_training = False
  
              if len(test_error_delta_buff) >= max_unimproving_models:
                  delta_mean = np.mean(
                      test_error_delta_buff[-max_unimproving_models:]
                  )
                  if delta_mean <= 0:
                      keep_training = False
              if callback is not None:
                  callback(test_error, real_buff, predicted_buff)
  
      return (
          self.best_model,
          lowest_test_error,
          int(time.time() - started),
      )
  
file_path: lightwood/lightwood/encoder/categorical/gym.py
imports: import copy
  import time
  import torch
  import numpy
  from lightwood.helpers.torch import LightwoodAutocast
  from lightwood.helpers.log import log
label_desc: Use image node_id 2 for calling the Gym obj's underlying member method code with example usage: obj.fit(train_data_loader, test_data_loader, desired_error, max_time, callback, eval_every_x_epochs, max_unimproving_models, custom_train_func, custom_test_func) and returns: self,  lowest_test_error,  int