code: def build_model(args):
      od = OrderedDict()
      for i in range(args.num_layers):
          if i == 0:
              od["linear_layer_%d" % (i + 1)] = torch.nn.Linear(
                  args.input_features, args.hidden_features
              )
              od["layer_norm_%d" % (i + 1)] = torch.nn.LayerNorm(
                  [args.batch_size, args.hidden_features]
              )
          elif i == args.num_layers - 1:
              od["linear_layer_%d" % (i + 1)] = torch.nn.Linear(
                  args.hidden_features, args.output_features
              )
              od["layer_norm_%d" % (i + 1)] = torch.nn.LayerNorm(
                  [args.batch_size, args.output_features]
              )
          else:
              od["linear_layer_%d" % (i + 1)] = torch.nn.Linear(
                  args.hidden_features, args.hidden_features
              )
              od["layer_norm_%d" % (i + 1)] = torch.nn.LayerNorm(
                  [args.batch_size, args.hidden_features]
              )
      return torch.nn.Sequential(od)
  
file_path: apex/apex/contrib/sparsity/test/checkpointing_test_reference.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 1 for calling a global function with example usage: build_model(args) and returns: torch