code: def train(
      args,
      extra_state: Dict[str, Any],
      trainer,
      task,
      epoch_itr,
      checkpoint_manager: Optional[checkpoint.CheckpointManager],
      output_queue: Optional[mp_queues.Queue] = None,
      **train_step_kwargs,
  ):
      # offset for current epoch (may be different from checkpoint offset)
      starting_offset = extra_state["batch_offset"]
  
      # Train until the learning rate gets too small
      max_epoch = args.max_epoch or math.inf
      lr = trainer.get_lr()
      train_meter = StopwatchMeter()
      train_meter.start()
      stop_training_mid_epoch = False
      stop_training_end_of_epoch = False
  
      do_prune = args.pruning_percentile > 0
      if do_prune:
          prune_masks = create_prune_masks(args, trainer)
          apply_prune_masks(prune_masks, trainer)
  
      while lr > args.min_lr and extra_state["epoch"] <= max_epoch:
          """Train the model for one epoch."""
  
          itr, progress, extra_meters = setup_epoch(
              args=args, epoch_itr=epoch_itr, trainer=trainer
          )
  
          for i, samples in enumerate(progress, start=starting_offset):
              clear_per_step_extra_state(extra_state)
              extra_state["num_iterations"] = (
                  extra_state.get("num_iterations", 0) + 1
              )
              if (
                  train_step_kwargs is not None
                  and "augment_adv" in train_step_kwargs.keys()
              ):
                  train_step_kwargs["augment_adv"] = (
                      extra_state["num_iterations"] > args.warmup_steps
                  )
              try:
                  log_output = trainer.train_step(
                      samples, **train_step_kwargs
                  )
              # Fairseq's fp16_trainer raises this uncommon error to indicate
              # that we should stop training.
              except FloatingPointError as e:
                  print(f"Stopping training due to: {e}.")
                  stop_training_mid_epoch = True
                  break
  
              if do_prune:
                  apply_prune_masks(prune_masks, trainer)
  
              if i == starting_offset:
                  # ignore the first mini-batch in words-per-second calculation
                  trainer.get_meter("wps").reset()
  
              # Clear any remaining metrics from previous steps. This should already
              # have been done before, but just in case - to make sure we catch
              # any case where extra_case does not get populated correctly.
              extra_state = clear_per_step_extra_state(extra_state)
              extra_state["batch_offset"] = i + 1
              (
                  extra_state,
                  stop_training_mid_epoch,
              ) = evals.save_and_eval(
                  args=args,
                  trainer=trainer,
                  task=task,
                  extra_state=extra_state,
                  checkpoint_manager=checkpoint_manager,
              )
  
              # This should come after save_and_eval. Even if log_output is None,
              # meaning that there was an overflow,  We should still run
              # save_and_eval to sync all_reduce and then skip the batch.
              if log_output is None:
                  # This indicates that the batch was skipped, typically
                  # because of OOM or FP16 overflow.
                  continue
  
              train_stats = evals.log_mid_epoch_stats(
                  trainer=trainer,
                  progress=progress,
                  extra_meters=extra_meters,
                  log_output=log_output,
              )
              extra_state = update_output(
                  args=args,
                  extra_state=extra_state,
                  output_queue=output_queue,
                  num_updates=trainer.get_num_updates(),
                  train_ppl=train_stats["ppl"],
                  # We only report wps at the end of an epoch, since
                  # the meter gets reset at the start of every epoch.
                  wps=None,
              )
  
              if (
                  hasattr(args, "lr_shrink")
                  and args.save_interval_updates > 0
                  and extra_state["num_iterations"]
                  % args.save_interval_updates
                  == 0
                  and args.shrink_lr_no_best_tune_loss > 0
                  and extra_state["tune_eval"]["num_since_best"]
                  > args.shrink_lr_no_best_tune_loss
              ):
                  current_lr = trainer.optimizer.get_lr()
                  trainer.optimizer.set_lr(current_lr * args.lr_shrink)
                  lr = trainer.optimizer.get_lr()
                  print(f"Decayed lr from {current_lr} to {lr}.")
  
              if stop_training_mid_epoch:
                  break
  
          # log end-of-epoch stats
          train_stats = evals.log_end_epoch_stats(
              trainer=trainer,
              progress=progress,
              extra_meters=extra_meters,
          )
  
          # batch_offset being None denotes the end of an epoch.
          extra_state["batch_offset"] = None
          extra_state, stop_training_end_of_epoch = evals.save_and_eval(
              args=args,
              trainer=trainer,
              task=task,
              extra_state=extra_state,
              end_of_epoch=True,
              checkpoint_manager=checkpoint_manager,
          )
          extra_state = update_output(
              args=args,
              extra_state=extra_state,
              output_queue=output_queue,
              num_updates=trainer.get_num_updates(),
              train_ppl=train_stats["ppl"],
              wps=train_stats["wps"],
          )
  
          if stop_training_mid_epoch or stop_training_end_of_epoch:
              break
  
          lr = trainer.lr_step(
              extra_state["epoch"], extra_state["tune_eval"]["loss"]
          )
          extra_state["epoch"] += 1
          extra_state["batch_offset"] = 0
          starting_offset = 0
  
      train_meter.stop()
      print(f"| done training in {train_meter.sum:.1f} seconds")
  
      # the checkpoint manager may be None
      if checkpoint_manager:
          checkpoint_manager.remove_all_checkpoints()
  
imports: import builtins
  import collections
  import datetime
  import faulthandler
  import math
  import multiprocessing.queues
  import os
  import queue
  import random
  import time
  from typing import Any, Callable, Dict, Optional, Tuple
  import numpy
  import torch
  from fairseq import data, distributed_utils, options, progress_bar, tasks, utils
  from fairseq.meters import AverageMeter, StopwatchMeter
  from fairseq.trainer import Trainer
  from pytorch_translate import char_aware_hybrid, char_source_hybrid, checkpoint, constants, evals, hybrid_transformer_rnn, multi_model, options, preprocess, sequence_criterions, transformer, transformer_aan, weighted_criterions
  from pytorch_translate.data import dictionary
  from pytorch_translate.dual_learning import dual_learning_criterion, dual_learning_task
  from pytorch_translate.file_io import PathManager
  from pytorch_translate.research.deliberation_networks import deliberation_networks
  from pytorch_translate.research.knowledge_distillation import dual_decoder_kd_loss, dual_decoder_kd_model, hybrid_dual_decoder_kd_model, knowledge_distillation_loss
  from pytorch_translate.research.rescore import cloze_transformer_model, rescoring_criterion
  from pytorch_translate.word_prediction import word_prediction_criterion, word_prediction_model
  from pytorch_translate import char_source_model, char_source_transformer_model, rnn, semi_supervised
label_desc: Use image node_id 16 for calling a global function with example usage: train(args, extra_state, trainer, task, epoch_itr, checkpoint_manager, output_queue) without return types