code: def main(step, args, model_state_dict, optimizer_state_dict):
      #
      # PART2
      #
  
      model = build_model(args).cuda()
      one_ll = next(model.children()).weight
      optimizer = FusedAdam(model.parameters())
      ASP.init_model_for_pruning(
          model,
          args.pattern,
          verbosity=args.verbosity,
          whitelist=args.whitelist,
          allow_recompute_mask=args.allow_recompute_mask,
      )
      ASP.init_optimizer_for_pruning(optimizer)
  
      torch.manual_seed(args.seed2)
      model.load_state_dict(model_state_dict)
      optimizer.load_state_dict(optimizer_state_dict)
  
      print(
          "Model sparsity is %s"
          % ("enabled" if ASP.is_sparsity_enabled() else "disabled")
      )
  
      # train for a few steps with sparse weights
      print("SPARSE :: ", one_ll)
      step = train_loop(
          args, model, optimizer, step, args.num_sparse_steps_2
      )
  
file_path: apex/apex/contrib/sparsity/test/checkpointing_test_part2.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 4 for calling a global function with example usage: main(step, args, model_state_dict, optimizer_state_dict) without return types