docs: def _train_model(
          self, train_data, loss_fn, valid_data=None, log_writer=None, restore_state={}
      ):
          """The internal training routine called by train_model() after setup
  
          Args:
              train_data: a tuple of Tensors (X,Y), a Dataset, or a DataLoader of
                  X (data) and Y (labels) for the train split
              loss_fn: the loss function to minimize (maps *data -> loss)
              valid_data: a tuple of Tensors (X,Y), a Dataset, or a DataLoader of
                  X (data) and Y (labels) for the dev split
              restore_state: a dictionary containing model weights (optimizer, main network) and training information
  
          If valid_data is not provided, then no checkpointing or
          evaluation on the dev set will occur.
          """
          # Set model to train mode
          self.train()
          train_config = self.config["train_config"]
  
          # Convert data to DataLoaders
          train_loader = self._create_data_loader(train_data)
          valid_loader = self._create_data_loader(valid_data)
          epoch_size = len(train_loader.dataset)
  
          # Move model to GPU
          if self.config["verbose"] and self.config["device"] != "cpu":
              print("Using GPU...")
          self.to(self.config["device"])
  
          # Set training components
          self._set_writer(train_config)
          self._set_logger(train_config, epoch_size)
          self._set_checkpointer(train_config)
          self._set_optimizer(train_config)
          self._set_scheduler(train_config)
  
          # Restore model if necessary
          if restore_state:
              start_iteration = self._restore_training_state(restore_state)
          else:
              start_iteration = 0
  
          # Train the model
          metrics_hist = {}  # The most recently seen value for all metrics
          for epoch in range(start_iteration, train_config["n_epochs"]):
              progress_bar = (
                  train_config["progress_bar"]
                  and self.config["verbose"]
                  and self.logger.log_unit == "epochs"
              )
  
              t = tqdm(
                  enumerate(train_loader),
                  total=len(train_loader),
                  disable=(not progress_bar),
              )
  
              self.running_loss = 0.0
              self.running_examples = 0
              for batch_num, data in t:
                  # NOTE: actual batch_size may not equal config's target batch_size
                  batch_size = len(data[0])
  
                  # Moving data to device
                  if self.config["device"] != "cpu":
                      data = place_on_gpu(data)
  
                  # Zero the parameter gradients
                  self.optimizer.zero_grad()
  
                  # Forward pass to calculate the average loss per example
                  loss = loss_fn(*data)
                  if torch.isnan(loss):
                      msg = "Loss is NaN. Consider reducing learning rate."
                      raise Exception(msg)
  
                  # Backward pass to calculate gradients
                  # Loss is an average loss per example
                  loss.backward()
  
                  # Perform optimizer step
                  self.optimizer.step()
  
                  # Calculate metrics, log, and checkpoint as necessary
                  metrics_dict = self._execute_logging(
                      train_loader, valid_loader, loss, batch_size
                  )
                  metrics_hist.update(metrics_dict)
  
                  # tqdm output
                  t.set_postfix(loss=metrics_dict["train/loss"])
  
              # Apply learning rate scheduler
              self._update_scheduler(epoch, metrics_hist)
  
          self.eval()
  
          # Restore best model if applicable
          if self.checkpointer:
              self.checkpointer.load_best_model(model=self)
  
          # Write log if applicable
          if self.writer:
              if self.writer.include_config:
                  self.writer.add_config(self.config)
              self.writer.close()
  
          # Print confusion matrix if applicable
          if self.config["verbose"]:
              print("Finished Training")
              if valid_loader is not None:
                  self.score(
                      valid_loader,
                      metric=train_config["validation_metric"],
                      verbose=True,
                      print_confusion_matrix=True,
                  )