code: def train(self, model, continued=False, start_epoch=0):
      logger.info("Config:")
      for name, value in self.config.items():
          logger.info("%s: %s" % (name, value))
      logger.info("Model parameters:")
      for name, param in model.named_parameters():
          logger.info("%s: %s" % (name, tuple(param.shape)))
  
      epochs, grad_accum = (
          self.config["num_epochs"],
          self.config["gradient_accumulation_steps"],
      )
      batch_size = self.config["batch_size"]
  
      # Set up data
      (
          examples_train,
          examples_dev,
          examples_test,
      ) = self.data.get_tensor_examples()
      stored_info = self.data.get_stored_info()
  
      # Set up optimizer and scheduler
      total_update_steps = (
          len(examples_train) * epochs // (grad_accum * batch_size)
      )
      if not continued:
          self.optimizer = self.get_optimizer(model)
          self.scheduler = self.get_scheduler(
              self.optimizer, total_update_steps
          )
  
      # Get model parameters for grad clipping
      plm_param, task_param = model.get_params()
  
      # Start training
      logger.info("*******************Training*******************")
      logger.info("Num samples: %d" % len(examples_train))
      logger.info("Num epochs: %d" % epochs)
      logger.info("Gradient accumulation steps: %d" % grad_accum)
      logger.info("Total update steps: %d" % total_update_steps)
      logger.info("Starting step: %d" % self.scheduler._step_count)
  
      loss_during_accum = []  # To compute effective loss at each update
      loss_during_report = 0.0  # Effective loss during logging step
      loss_history = (
          []
      )  # Full history of effective loss; length equals total update steps
      max_f1, max_f1_test = 0, 0
      start_time = time.time()
      if type(self.optimizer) == FusedAdam:
          self.optimizer.zero_grad()
      else:
          self.optimizer.zero_grad(set_to_none=True)
  
      trainloader = DataLoader(
          examples_train,
          batch_size=batch_size,
          shuffle=True,
          num_workers=0,
          collate_fn=self.collate_fn,
          pin_memory=True,
      )
  
      for epo in range(start_epoch, epochs):
          logger.info(
              "*******************EPOCH %d*******************" % epo
          )
          for doc_key, example in trainloader:
              # Forward pass
              model.train()
              example_gpu = {}
              for k, v in example.items():
                  if v is not None:
                      example_gpu[k] = v.to(self.device)
  
              with torch.cuda.amp.autocast(
                  enabled=self.use_amp, dtype=torch.bfloat16
              ):
                  # example_gpu['lr_pair_flag']=torch.rand([1, 343, 11, 4]).to(self.device)
                  # print(example_gpu)
                  loss = model(**example_gpu) / grad_accum
              # Backward; accumulate gradients and clip by grad norm
              loss.backward()
              loss_during_accum.append(loss.item())
              # Update
              if len(loss_during_accum) % grad_accum == 0:
                  if self.config["max_grad_norm"]:
                      norm_plm = torch.nn.utils.clip_grad_norm_(
                          plm_param,
                          self.config["max_grad_norm"],
                          error_if_nonfinite=False,
                      )
                      norm_task = torch.nn.utils.clip_grad_norm_(
                          task_param,
                          self.config["max_grad_norm"],
                          error_if_nonfinite=False,
                      )
                  self.optimizer.step()
                  self.scheduler.step()
                  if type(self.optimizer) == FusedAdam:
                      self.optimizer.zero_grad()
                  else:
                      self.optimizer.zero_grad(set_to_none=True)
  
                  # Compute effective loss
                  effective_loss = np.sum(loss_during_accum).item()
                  loss_during_accum = []
                  loss_during_report += effective_loss
                  loss_history.append(effective_loss)
  
                  # Report
                  if (
                      self.scheduler._step_count
                      % self.config["report_frequency"]
                      == 0
                  ):
                      # Show avg loss during last report interval
                      avg_loss = (
                          loss_during_report
                          / self.config["report_frequency"]
                      )
                      loss_during_report = 0.0
                      end_time = time.time()
                      logger.info(
                          "Step %d: avg loss %.2f; steps/sec %.2f"
                          % (
                              self.scheduler._step_count,
                              avg_loss,
                              self.config["report_frequency"]
                              / (end_time - start_time),
                          )
                      )
                      start_time = end_time
                      wandb.log({"avg_loss": avg_loss})
  
                  # Evaluate
                  if (
                      self.scheduler._step_count
                      % self.config["eval_frequency"]
                      == 0
                  ):
                      logger.info("Dev")
  
                      f1, _ = self.evaluate(
                          model,
                          examples_dev,
                          stored_info,
                          self.scheduler._step_count,
                      )
                      logger.info("Test")
                      f1_test = 0.0
                      if f1 > max_f1:
                          max_f1 = max(max_f1, f1)
                          max_f1_test = 0.0
                          self.save_model_checkpoint(
                              model,
                              self.optimizer,
                              self.scheduler,
                              self.scheduler._step_count,
                              epo,
                          )
  
                      logger.info("Eval max f1: %.2f" % max_f1)
                      logger.info("Test max f1: %.2f" % max_f1_test)
                      start_time = time.time()
  
      logger.info("**********Finished training**********")
      logger.info(
          "Actual update steps: %d" % self.scheduler._step_count
      )
  
      return
  
imports: import sys
  import os
  import logging
  import random
  import numpy
  import torch
  from torch import nn
  from torch.optim import AdamW
  from util.multigpu_fused_adam import FusedAdam
  from torch.optim.lr_scheduler import LambdaLR
  from torch.utils.data import DataLoader
  import time
  from os.path import join
  from datetime import datetime
  import util
  from util.tensorize_coref import CorefDataProcessor, coref_collate_fn
  from util.tensorize_ner import NERDataProcessor, ner_collate_fn
  from util.tensorize_ere import EREDataProcessor, ere_collate_fn
  from metrics import CorefEvaluator, MentionEvaluator
  from models.model_coref import CorefModel
  from models.model_ner import NERModel
  from models.model_ere import EREModel
  import wandb
label_desc: Use image node_id 3 for calling the Runner obj's underlying member method code with example usage: obj.train(model, continued, start_epoch) without return types