code: def train_step(
      args, model, optimizer, input_batch, target_batch, step
  ):
      predicted_target = model(input_batch)
      loss = ((predicted_target - target_batch) ** 2).sum()
      loss.backward()
      optimizer.step()
      optimizer.zero_grad()
      step = step + 1
      # print("Step %d :: loss=%e" % (step, loss.item()))
      return step
  
file_path: apex/apex/contrib/sparsity/test/toy_problem.py
imports: from collections import OrderedDict
  import torch
  from apex.optimizers import FusedAdam
  from apex.contrib.sparsity import ASP
label_desc: Use image node_id 2 for calling a global function with example usage: train_step(args, model, optimizer, input_batch, target_batch, step) and returns: step