class_docstr: A trainer for doing supervised learning with gradient descent. It just takes a labeled dataset
  and a `DataLoader`, and uses the supplied `Optimizer` to learn the weights for your model over
  some fixed number of epochs. You can also pass in a validation data_loader and enable early
  stopping. There are many other bells and whistles as well.
  
  Registered as a `Trainer` with the name "gradient_descent" (and is also the default `Trainer`).
  The constructor that is registered is [`from_partial_objects`](#from_partial_objects) -
  see the arguments to that function for the exact keys that should be used, if you are using
  a configuration file. They largely match the arguments to `__init__`, and we don't repeat their
  docstrings in `from_partial_objects`.
  
  [0]: https://tinyurl.com/y5mv44fw
  
  # Parameters
  
  model : `Model`, required.
      An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
      their `forward` method returns a dictionary with a "loss" key, containing a
      scalar tensor representing the loss function to be optimized.
  
      If you are training your model using GPUs, your model should already be
      on the correct device. (If you are using our `train` command this will be
      handled for you.)
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  optimizer : `torch.nn.Optimizer`, required.
      An instance of a Pytorch Optimizer, instantiated with the parameters of the
      model to be optimized.
  
  data_loader : `DataLoader`, required.
      A `DataLoader` containing your `Dataset`, yielding padded indexed batches.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  patience : `Optional[int] > 0`, optional (default=`None`)
      Number of epochs to be patient before early stopping: the training is stopped
      after `patience` epochs with no improvement. If given, it must be `> 0`.
      If None, early stopping is disabled.
  
  validation_metric : `Union[str, List[str]]`, optional (default=`"-loss"`)
      Validation metric to measure for whether to stop training using patience
      and whether to serialize an `is_best` model each epoch. The metric name
      must be prepended with either "+" or "-", which specifies whether the metric
      is an increasing or decreasing function. If you specify more than one metric,
      the metrics will be summed to make the `is_best` decision.
  
  validation_data_loader : `DataLoader`, optional (default=`None`)
      A `DataLoader` to use for the validation set.  If `None`, then
      use the training `DataLoader` with the validation data.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  num_epochs : `int`, optional (default = `20`)
      Number of training epochs.
  
  serialization_dir : `str`, optional (default=`None`)
      Path to directory for saving and loading model files. Models will not be saved if
      this parameter is not passed.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  checkpointer : `Checkpointer`, optional (default=`None`)
      A `Checkpointer` is responsible for periodically saving model weights.  If none is given
      here, we will construct one with default parameters.
  
  cuda_device : `Optional[Union[int, torch.device]]`, optional (default = `None`)
      An integer or `torch.device` specifying the CUDA device to use for this process.
      If -1, the CPU is used. If `None` and you have a GPU available, that GPU will be used.
  
      !!! Note
          If you *don't* intend to use a GPU, but you have one available, you'll need
          to explicitly set `cuda_device=-1`.
  
      !!! Note
          If you intend to use a GPU, your model already needs to be on the correct device,
          which you can do with `model = model.cuda()`.
  
      !!! Note
          Data parallelism is controlled at the allennlp train level, so each trainer will have a single GPU.
  
  grad_norm : `Union[float, bool]`, optional (default = `False`)
      If a float, gradient norms will be rescaled to have a maximum of this value.
      If `True`, the gradient norms will be calculated and passed through to any `TrainerCallbacks`,
      but won't be rescaled.
      If `False`, gradient norms will not be calculated or rescaled.
  
  grad_clipping : `float`, optional (default = `None`)
      If provided, gradients will be clipped `during the backward pass` to have an (absolute)
      maximum of this value.  If you are getting `NaNs` in your gradients during training
      that are not solved by using `grad_norm`, you may need this.
  
  learning_rate_scheduler : `LearningRateScheduler`, optional (default = `None`)
      If specified, the learning rate will be decayed with respect to
      this schedule at the end of each epoch (or batch, if the scheduler implements
      the `step_batch` method). If you use `torch.optim.lr_scheduler.ReduceLROnPlateau`,
      this will use the `validation_metric` provided to determine if learning has plateaued.
      To support updating the learning rate on every batch, this can optionally implement
      `step_batch(batch_num_total)` which updates the learning rate given the batch number.
  
  momentum_scheduler : `MomentumScheduler`, optional (default = `None`)
      If specified, the momentum will be updated at the end of each batch or epoch
      according to the schedule.
  
  moving_average : `MovingAverage`, optional, (default = `None`)
      If provided, we will maintain moving averages for all parameters. During training, we
      employ a shadow variable for each parameter, which maintains the moving average. During
      evaluation, we backup the original parameters and assign the moving averages to corresponding
      parameters. Be careful that when saving the checkpoint, we will save the moving averages of
      parameters. This is necessary because we want the saved model to perform as well as the validated
      model if we load it later. But this may cause problems if you restart the training from checkpoint.
  
  callbacks : `List[TrainerCallback]`, optional (default = `None`)
      A list of callbacks that can be called at certain events: e.g. each batch, epoch, and at the start
      and end of training, etc.
  
  distributed : `bool`, optional, (default = `False`)
      If set, PyTorch's `DistributedDataParallel` is used to train the model in multiple GPUs. This also
      requires `world_size` to be greater than 1.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately (you need a top-level "distributed" key, next to
      the "trainer" entry, that specifies a list of "cuda_devices").
  
  local_rank : `int`, optional, (default = `0`)
      This is the unique identifier of the `Trainer` in a distributed process group. The GPU device id is
      used as the rank.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  world_size : `int`, (default = `1`)
      The number of `Trainer` workers participating in the distributed training.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  num_gradient_accumulation_steps : `int`, optional, (default = `1`)
      Gradients are accumulated for the given number of steps before doing an optimizer step. This can
      be useful to accommodate batches that are larger than the RAM size. Refer [Thomas Wolf's
      post][0] for details on Gradient Accumulation.
  
  use_amp : `bool`, optional, (default = `False`)
      If `True`, we'll train using [Automatic Mixed Precision](https://pytorch.org/docs/stable/amp.html).
  
  enable_default_callbacks : `bool`, optional (default = `True`)
      When `True`, the [`DEFAULT_CALLBACKS`](#default_callbacks) will be used in
      addition to any other callbacks listed in the `callbacks` parameter.
      When set to `False`, `DEFAULT_CALLBACKS` are not used.
  
  run_confidence_checks : `bool`, optional (default = `True`)
      Determines whether model confidence checks, such as
      [`NormalizationBiasVerification`](../../confidence_checks/normalization_bias_verification/),
      are run.
  
  run_sanity_checks : `bool`, optional (default = `True`)
      This parameter is deprecated. Please use `run_confidence_checks` instead.
  
  grad_scaling : `bool`, optional (default = `True`)
      When `use_amp` is `True`, this determines whether or not to use a [`GradScaler`]
      (https://pytorch.org/docs/stable/amp.html?highlight=gradscaler#torch.cuda.amp.GradScaler).
  
      !!! Note
          This parameter is ignored when `use_amp` is `False`.
  
  ddp_wrapped_model : `Optional[DdpWrappedModel]`, optional (default = `None`)
      The `model` wrapped with a `DdpAccelerator` for distributed training.
  
      !!! Note
          This is required for distributed training.
class_docstr_tok: ["A","trainer","for","doing","supervised","learning","with","gradient","descent",".","It","just","takes","a","labeled","dataset","and","a","`","DataLoader","`",",","and","uses","the","supplied","`","Optimizer","`","to","learn","the","weights","for","your","model","over","some","fixed","number","of","epochs",".","You","can","also","pass","in","a","validation","data_loader","and","enable","early","stopping",".","There","are","many","other","bells","and","whistles","as","well",".","Registered","as","a","`","Trainer","`","with","the","name","``","gradient_descent","''","(","and","is","also","the","default","`","Trainer","`",")",".","The","constructor","that","is","registered","is","[","`","from_partial_objects","`","]","(","#","from_partial_objects",")","-","see","the","arguments","to","that","function","for","the","exact","keys","that","should","be","used",",","if","you","are","using","a","configuration","file",".","They","largely","match","the","arguments","to","`","__init__","`",",","and","we","do","n't","repeat","their","docstrings","in","`","from_partial_objects","`",".","[","0","]",":","https",":","\/\/tinyurl.com\/y5mv44fw","#","Parameters","model",":","`","Model","`",",","required",".","An","AllenNLP","model","to","be","optimized",".","Pytorch","Modules","can","also","be","optimized","if","their","`","forward","`","method","returns","a","dictionary","with","a","``","loss","''","key",",","containing","a","scalar","tensor","representing","the","loss","function","to","be","optimized",".","If","you","are","training","your","model","using","GPUs",",","your","model","should","already","be","on","the","correct","device",".","(","If","you","are","using","our","`","train","`","command","this","will","be","handled","for","you",".",")","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","optimizer",":","`","torch.nn.Optimizer","`",",","required",".","An","instance","of","a","Pytorch","Optimizer",",","instantiated","with","the","parameters","of","the","model","to","be","optimized",".","data_loader",":","`","DataLoader","`",",","required",".","A","`","DataLoader","`","containing","your","`","Dataset","`",",","yielding","padded","indexed","batches",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","patience",":","`","Optional","[","int","]",">","0","`",",","optional","(","default=","`","None","`",")","Number","of","epochs","to","be","patient","before","early","stopping",":","the","training","is","stopped","after","`","patience","`","epochs","with","no","improvement",".","If","given",",","it","must","be","`",">","0","`",".","If","None",",","early","stopping","is","disabled",".","validation_metric",":","`","Union","[","str",",","List","[","str","]","]","`",",","optional","(","default=","`","``","-loss","''","`",")","Validation","metric","to","measure","for","whether","to","stop","training","using","patience","and","whether","to","serialize","an","`","is_best","`","model","each","epoch",".","The","metric","name","must","be","prepended","with","either","``","+","''","or","``","-","''",",","which","specifies","whether","the","metric","is","an","increasing","or","decreasing","function",".","If","you","specify","more","than","one","metric",",","the","metrics","will","be","summed","to","make","the","`","is_best","`","decision",".","validation_data_loader",":","`","DataLoader","`",",","optional","(","default=","`","None","`",")","A","`","DataLoader","`","to","use","for","the","validation","set",".","If","`","None","`",",","then","use","the","training","`","DataLoader","`","with","the","validation","data",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","num_epochs",":","`","int","`",",","optional","(","default","=","`","20","`",")","Number","of","training","epochs",".","serialization_dir",":","`","str","`",",","optional","(","default=","`","None","`",")","Path","to","directory","for","saving","and","loading","model","files",".","Models","will","not","be","saved","if","this","parameter","is","not","passed",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","checkpointer",":","`","Checkpointer","`",",","optional","(","default=","`","None","`",")","A","`","Checkpointer","`","is","responsible","for","periodically","saving","model","weights",".","If","none","is","given","here",",","we","will","construct","one","with","default","parameters",".","cuda_device",":","`","Optional","[","Union","[","int",",","torch.device","]","]","`",",","optional","(","default","=","`","None","`",")","An","integer","or","`","torch.device","`","specifying","the","CUDA","device","to","use","for","this","process",".","If","-1",",","the","CPU","is","used",".","If","`","None","`","and","you","have","a","GPU","available",",","that","GPU","will","be","used",".","!","!","!","Note","If","you","*","do","n't","*","intend","to","use","a","GPU",",","but","you","have","one","available",",","you","'ll","need","to","explicitly","set","`","cuda_device=-1","`",".","!","!","!","Note","If","you","intend","to","use","a","GPU",",","your","model","already","needs","to","be","on","the","correct","device",",","which","you","can","do","with","`","model","=","model.cuda","(",")","`",".","!","!","!","Note","Data","parallelism","is","controlled","at","the","allennlp","train","level",",","so","each","trainer","will","have","a","single","GPU",".","grad_norm",":","`","Union","[","float",",","bool","]","`",",","optional","(","default","=","`","False","`",")","If","a","float",",","gradient","norms","will","be","rescaled","to","have","a","maximum","of","this","value",".","If","`","True","`",",","the","gradient","norms","will","be","calculated","and","passed","through","to","any","`","TrainerCallbacks","`",",","but","wo","n't","be","rescaled",".","If","`","False","`",",","gradient","norms","will","not","be","calculated","or","rescaled",".","grad_clipping",":","`","float","`",",","optional","(","default","=","`","None","`",")","If","provided",",","gradients","will","be","clipped","`","during","the","backward","pass","`","to","have","an","(","absolute",")","maximum","of","this","value",".","If","you","are","getting","`","NaNs","`","in","your","gradients","during","training","that","are","not","solved","by","using","`","grad_norm","`",",","you","may","need","this",".","learning_rate_scheduler",":","`","LearningRateScheduler","`",",","optional","(","default","=","`","None","`",")","If","specified",",","the","learning","rate","will","be","decayed","with","respect","to","this","schedule","at","the","end","of","each","epoch","(","or","batch",",","if","the","scheduler","implements","the","`","step_batch","`","method",")",".","If","you","use","`","torch.optim.lr_scheduler.ReduceLROnPlateau","`",",","this","will","use","the","`","validation_metric","`","provided","to","determine","if","learning","has","plateaued",".","To","support","updating","the","learning","rate","on","every","batch",",","this","can","optionally","implement","`","step_batch","(","batch_num_total",")","`","which","updates","the","learning","rate","given","the","batch","number",".","momentum_scheduler",":","`","MomentumScheduler","`",",","optional","(","default","=","`","None","`",")","If","specified",",","the","momentum","will","be","updated","at","the","end","of","each","batch","or","epoch","according","to","the","schedule",".","moving_average",":","`","MovingAverage","`",",","optional",",","(","default","=","`","None","`",")","If","provided",",","we","will","maintain","moving","averages","for","all","parameters",".","During","training",",","we","employ","a","shadow","variable","for","each","parameter",",","which","maintains","the","moving","average",".","During","evaluation",",","we","backup","the","original","parameters","and","assign","the","moving","averages","to","corresponding","parameters",".","Be","careful","that","when","saving","the","checkpoint",",","we","will","save","the","moving","averages","of","parameters",".","This","is","necessary","because","we","want","the","saved","model","to","perform","as","well","as","the","validated","model","if","we","load","it","later",".","But","this","may","cause","problems","if","you","restart","the","training","from","checkpoint",".","callbacks",":","`","List","[","TrainerCallback","]","`",",","optional","(","default","=","`","None","`",")","A","list","of","callbacks","that","can","be","called","at","certain","events",":","e.g",".","each","batch",",","epoch",",","and","at","the","start","and","end","of","training",",","etc",".","distributed",":","`","bool","`",",","optional",",","(","default","=","`","False","`",")","If","set",",","PyTorch","'s","`","DistributedDataParallel","`","is","used","to","train","the","model","in","multiple","GPUs",".","This","also","requires","`","world_size","`","to","be","greater","than","1",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately","(","you","need","a","top-level","``","distributed","''","key",",","next","to","the","``","trainer","''","entry",",","that","specifies","a","list","of","``","cuda_devices","''",")",".","local_rank",":","`","int","`",",","optional",",","(","default","=","`","0","`",")","This","is","the","unique","identifier","of","the","`","Trainer","`","in","a","distributed","process","group",".","The","GPU","device","id","is","used","as","the","rank",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","world_size",":","`","int","`",",","(","default","=","`","1","`",")","The","number","of","`","Trainer","`","workers","participating","in","the","distributed","training",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","num_gradient_accumulation_steps",":","`","int","`",",","optional",",","(","default","=","`","1","`",")","Gradients","are","accumulated","for","the","given","number","of","steps","before","doing","an","optimizer","step",".","This","can","be","useful","to","accommodate","batches","that","are","larger","than","the","RAM","size",".","Refer","[","Thomas","Wolf's","post","]","[","0","]","for","details","on","Gradient","Accumulation",".","use_amp",":","`","bool","`",",","optional",",","(","default","=","`","False","`",")","If","`","True","`",",","we","'ll","train","using","[","Automatic","Mixed","Precision","]","(","https",":","\/\/pytorch.org\/docs\/stable\/amp.html",")",".","enable_default_callbacks",":","`","bool","`",",","optional","(","default","=","`","True","`",")","When","`","True","`",",","the","[","`","DEFAULT_CALLBACKS","`","]","(","#","default_callbacks",")","will","be","used","in","addition","to","any","other","callbacks","listed","in","the","`","callbacks","`","parameter",".","When","set","to","`","False","`",",","`","DEFAULT_CALLBACKS","`","are","not","used",".","run_confidence_checks",":","`","bool","`",",","optional","(","default","=","`","True","`",")","Determines","whether","model","confidence","checks",",","such","as","[","`","NormalizationBiasVerification","`","]","(","..","\/","..","\/confidence_checks\/normalization_bias_verification\/",")",",","are","run",".","run_sanity_checks",":","`","bool","`",",","optional","(","default","=","`","True","`",")","This","parameter","is","deprecated",".","Please","use","`","run_confidence_checks","`","instead",".","grad_scaling",":","`","bool","`",",","optional","(","default","=","`","True","`",")","When","`","use_amp","`","is","`","True","`",",","this","determines","whether","or","not","to","use","a","[","`","GradScaler","`","]","(","https",":","\/\/pytorch.org\/docs\/stable\/amp.html","?","highlight=gradscaler","#","torch.cuda.amp.GradScaler",")",".","!","!","!","Note","This","parameter","is","ignored","when","`","use_amp","`","is","`","False","`",".","ddp_wrapped_model",":","`","Optional","[","DdpWrappedModel","]","`",",","optional","(","default","=","`","None","`",")","The","`","model","`","wrapped","with","a","`","DdpAccelerator","`","for","distributed","training",".","!","!","!","Note","This","is","required","for","distributed","training","."]
code: def _try_train(self) -> Tuple[Dict[str, Any], int]:
      logger.info("Beginning training.")
  
      val_metrics: Dict[str, float] = {}
      metrics: Dict[str, Any] = {}
      training_start_time = None
  
      metrics["best_epoch"] = self._metric_tracker.best_epoch
      for key, value in self._metric_tracker.best_epoch_metrics.items():
          metrics["best_validation_" + key] = value
  
      for epoch in range(self._num_epochs):
          epoch_start_time = time.time()
          train_metrics = self._train_epoch(epoch)
  
          if (
              self._epochs_completed
              < self._start_after_epochs_completed
          ):
              # We're still catching up with the checkpoint, so we do nothing.
              # Note that we have to call _train_epoch() even when we know the epoch is skipped. We have to
              # read from the data loader, because the data loader and dataset readers might use randomness,
              # and we have to make sure we consume exactly the same instances in exactly the same way every
              # time we train, even when starting from a checkpoint, so that we update the randomness
              # generators in the same way each time.
              self._epochs_completed += 1
              self._batches_in_epoch_completed = 0
              continue
          if training_start_time is None:
              training_start_time = epoch_start_time
  
          # get peak of memory usage
          for key, value in train_metrics.items():
              if key.startswith("gpu_") and key.endswith("_memory_MB"):
                  metrics["peak_" + key] = max(
                      metrics.get("peak_" + key, 0), value
                  )
              elif key.startswith("worker_") and key.endswith(
                  "_memory_MB"
              ):
                  metrics["peak_" + key] = max(
                      metrics.get("peak_" + key, 0), value
                  )
  
          this_epoch_val_metric: Optional[float] = None
          if (
              self._should_validate_this_epoch
              and self._validation_data_loader is not None
          ):
              with torch.no_grad():
                  # We have a validation set, so compute all the metrics on it.
                  (
                      val_loss,
                      val_reg_loss,
                      num_batches,
                  ) = self._validation_loss(epoch)
  
                  # It is safe again to wait till the validation is done. This is
                  # important to get the metrics right.
                  if self._distributed:
                      dist.barrier()
  
                  val_loss = dist_reduce_sum(val_loss)
                  num_batches = dist_reduce_sum(num_batches)
                  if val_reg_loss is not None:
                      val_reg_loss = dist_reduce_sum(val_reg_loss)
  
                  val_metrics = training_util.get_metrics(
                      self.model,
                      val_loss,
                      val_reg_loss,
                      batch_loss=None,
                      batch_reg_loss=None,
                      num_batches=num_batches,
                      reset=True,
                  )
  
                  # Check validation metric for early stopping
                  this_epoch_val_metric = (
                      self._metric_tracker.combined_score(val_metrics)
                  )
                  self._metric_tracker.add_metrics(val_metrics)
  
          # Create overall metrics dict
          training_elapsed_time = time.time() - training_start_time
          metrics["training_duration"] = str(
              datetime.timedelta(seconds=training_elapsed_time)
          )
          metrics["epoch"] = epoch
  
          for key, value in train_metrics.items():
              metrics["training_" + key] = value
          for key, value in val_metrics.items():
              metrics["validation_" + key] = value
  
          if (
              self._should_validate_this_epoch
              and self._metric_tracker.is_best_so_far()
          ):
              # Update all the best_ metrics.
              # (Otherwise they just stay the same as they were.)
              metrics["best_epoch"] = epoch
              for key, value in val_metrics.items():
                  metrics["best_validation_" + key] = value
  
              self._metric_tracker.best_epoch_metrics = val_metrics
  
          if self._serialization_dir and self._primary:
              common_util.dump_metrics(
                  os.path.join(
                      self._serialization_dir,
                      f"metrics_epoch_{epoch}.json",
                  ),
                  metrics,
              )
  
          # The Scheduler API is agnostic to whether your schedule requires a validation metric -
          # if it doesn't, the validation metric passed here is ignored.
          if self._learning_rate_scheduler:
              self._learning_rate_scheduler.step(this_epoch_val_metric)
          if self._momentum_scheduler:
              self._momentum_scheduler.step(this_epoch_val_metric)
          for callback in self._callbacks:
              callback.on_epoch(
                  self,
                  metrics=metrics,
                  epoch=epoch,
                  is_primary=self._primary,
              )
  
          self._epochs_completed += 1
          self._batches_in_epoch_completed = 0
  
          checkpoint_saved = False
          if self._checkpointer is not None:
              # The checkpointer saves state from the learning rate scheduler, momentum scheduler, moving
              # average, and callbacks, so we have to make sure those are updated before we save the
              # checkpoint here.
              checkpoint_saved = (
                  self._checkpointer.maybe_save_checkpoint(
                      self,
                      self._epochs_completed,
                      self._batches_in_epoch_completed,
                  )
              )
  
              # Wait for each primary process to finish saving the checkpoint
              if self._distributed:
                  dist.barrier()
  
          if (
              self._should_validate_this_epoch
              and self._serialization_dir
              and self._metric_tracker.is_best_so_far()
          ):
              should_save_model_state: bool
              if (
                  self._ddp_wrapped_model is not None
                  and self._ddp_wrapped_model.is_sharded
              ):
                  # Each worker saves its own shard for now (we combine the shards later).
                  self._best_model_filename = os.path.join(
                      self._serialization_dir, f"best_w{self._rank}.th"
                  )
                  should_save_model_state = True
              else:
                  self._best_model_filename = os.path.join(
                      self._serialization_dir, "best.th"
                  )
                  should_save_model_state = self._primary
  
              if should_save_model_state:
                  if self._moving_average is None:
                      # If we're not using a moving average and the checkpointer just saved a checkpoint,
                      # we can just copy over that model state checkpoint to the '_best_model_filename'.
                      # Otherwise we need to save the model state on our own.
                      if (
                          self._checkpointer is not None
                          and checkpoint_saved
                      ):
                          last_checkpoint = (
                              self._checkpointer.find_latest_checkpoint()
                          )
                          assert last_checkpoint is not None
                          model_state_file, _ = last_checkpoint
                          if os.path.exists(self._best_model_filename):
                              os.remove(self._best_model_filename)
                          hardlink_or_copy(
                              model_state_file,
                              self._best_model_filename,
                          )
                      else:
                          self._save_model_state(
                              self._best_model_filename
                          )
                  else:
                      self._moving_average.assign_average_value()
                      try:
                          self._save_model_state(
                              self._best_model_filename
                          )
                      finally:
                          self._moving_average.restore()
  
          # Wait for the primary process to finish saving the best
          if self._distributed:
              dist.barrier()
  
          epoch_elapsed_time = time.time() - epoch_start_time
          logger.info(
              "Epoch duration: %s",
              datetime.timedelta(seconds=epoch_elapsed_time),
          )
  
          if self._metric_tracker.should_stop_early():
              logger.info("Ran out of patience. Stopping training.")
              break
  
          if epoch < self._num_epochs - 1:
              time_per_epoch = training_elapsed_time / (
                  (epoch + 1) - self._start_after_epochs_completed
              )
              # Note: If the first non-skipped epoch is half skipped (because it was checkpointed half-way
              # through), then this estimate is going to be optimistic.
              estimated_time_remaining = (
                  time_per_epoch * self._num_epochs
              ) - training_elapsed_time
              formatted_time = str(
                  datetime.timedelta(
                      seconds=int(estimated_time_remaining)
                  )
              )
              logger.info(
                  "Estimated training time remaining: %s",
                  formatted_time,
              )
      else:
          epoch = self._num_epochs - 1
  
      # Load the best model state before returning
      if (
          self._best_model_filename is None
          or self._metric_tracker.is_best_so_far()
      ):
          self._finalize_model()
      else:
          # The model we're loading here has already been finalized.
          self._load_model_state(self._best_model_filename)
  
      return metrics, epoch
  
imports: import datetime
  import glob
  import logging
  import math
  import os
  import re
  import time
  import warnings
  from typing import Optional, Union, List, Dict, Tuple, Any, Type
  import torch
  from torch.cuda import amp
  from torch.nn.utils import clip_grad_norm_
  import torch.distributed
  from torch.cuda.amp.grad_scaler import OptState
  from allennlp.common.checks import ConfigurationError, check_for_gpu
  from allennlp.common import util, Tqdm, Lazy
  from allennlp.common.file_utils import hardlink_or_copy
  from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict
  from allennlp.models.model import Model
  from allennlp.nn.parallel import DdpAccelerator, DdpWrappedModel, TorchDdpAccelerator
  from allennlp.nn.util import dist_reduce_sum
  from allennlp.training.callbacks import ConsoleLoggerCallback
  from allennlp.training.callbacks.confidence_checks import ConfidenceChecksCallback
  from allennlp.training.callbacks.backward import MixedPrecisionBackwardCallback
  from allennlp.training.checkpointer import Checkpointer
  from allennlp.training.learning_rate_schedulers.learning_rate_scheduler import LearningRateScheduler
  from allennlp.training.metric_tracker import MetricTracker
  from allennlp.training.momentum_schedulers.momentum_scheduler import MomentumScheduler
  from allennlp.training.moving_average import MovingAverage
  from allennlp.training.optimizers import Optimizer
  from allennlp.training.trainer import Trainer, TrainerCheckpoint
  from allennlp.training.callbacks import TrainerCallback
  from allennlp.training import util
label_desc: Use image node_id 9 for calling the GradientDescentTrainer obj's underlying member method code with example usage: obj._try_train() and returns: metrics,  epoch