docstr: Normalization on the loss (gradient accumulation or distributed training), followed by
  backward pass on the loss, possibly followed by a parameter update (depending on the gradient accumulation).
  Also update the metrics for tensorboard.
docstr_tok: ["Normalization","on","the","loss","(","gradient","accumulation","or","distributed","training",")",",","followed","by","backward","pass","on","the","loss",",","possibly","followed","by","a","parameter","update","(","depending","on","the","gradient","accumulation",")",".","Also","update","the","metrics","for","tensorboard","."]
code: def optimize(self, loss):
      """
      Normalization on the loss (gradient accumulation or distributed training), followed by
      backward pass on the loss, possibly followed by a parameter update (depending on the gradient accumulation).
      Also update the metrics for tensorboard.
      """
      # Check for NaN
      if (loss != loss).data.any():
          logger.error("NaN detected")
          exit()
  
      if self.multi_gpu:
          loss = loss.mean()
      if self.params.gradient_accumulation_steps > 1:
          loss = loss / self.params.gradient_accumulation_steps
  
      if self.fp16:
          from apex import amp
  
          with amp.scale_loss(loss, self.optimizer) as scaled_loss:
              scaled_loss.backward()
      else:
          loss.backward()
  
      self.iter()
      if self.n_iter % self.params.gradient_accumulation_steps == 0:
          if self.fp16:
              nn.utils.clip_grad_norm_(
                  amp.master_params(self.optimizer),
                  self.params.max_grad_norm,
              )
          else:
              nn.utils.clip_grad_norm_(
                  self.student.parameters(), self.params.max_grad_norm
              )
          self.optimizer.step()
          self.optimizer.zero_grad()
          self.scheduler.step()
  
imports: import math
  import os
  import time
  import psutil
  import torch
  from grouped_batch_sampler import GroupedBatchSampler, create_lengths_groups
  from lm_seqs_dataset import LmSeqsDataset
  from torch import nn
  from torch.optim import AdamW
  from torch.utils.data import BatchSampler, DataLoader, RandomSampler
  from torch.utils.data.distributed import DistributedSampler
  from tqdm import tqdm
  from transformers import get_linear_schedule_with_warmup
  from utils import logger
label_desc: Use image node_id 7 for calling the Distiller obj's underlying member method code with example usage: obj.optimize(loss) without return types