code: def main(args):
      model = build_model(args).cuda()
      one_ll = next(model.children()).weight
      optimizer = FusedAdam(model.parameters())
      # only prune linear layers, even though we also support conv1d, conv2d and conv3d
      ASP.init_model_for_pruning(
          model,
          "m4n2_1d",
          whitelist=[torch.nn.Linear],
          allow_recompute_mask=True,
      )
      ASP.init_optimizer_for_pruning(optimizer)
  
      step = 0
  
      # train for a few steps with dense weights
      print("DENSE :: ", one_ll)
      step = train_loop(
          args, model, optimizer, step, args.num_dense_steps
      )
  
      # simulate sparsity by inserting zeros into existing dense weights
      ASP.compute_sparse_masks()
  
      # train for a few steps with sparse weights
      print("SPARSE :: ", one_ll)
      step = train_loop(
          args, model, optimizer, step, args.num_sparse_steps
      )
  
      # recompute sparse masks
      ASP.compute_sparse_masks()
  
      # train for a few steps with sparse weights
      print("SPARSE :: ", one_ll)
      step = train_loop(
          args, model, optimizer, step, args.num_sparse_steps_2
      )
  
      # turn off sparsity
      print("SPARSE :: ", one_ll)
      ASP.restore_pruned_weights()
  
      # train for a few steps with dense weights
      print("DENSE :: ", one_ll)
      step = train_loop(
          args, model, optimizer, step, args.num_dense_steps_2
      )
  
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 4 for calling a global function with example usage: main(args) without return types