code: def init_model():
      global cpu
  
      # MOE
      if args.model == "moe":
          model = MoENetwork(
              args.input_width,
              args.hidden_width,
              args.output_width,
              2**args.depth,
              1,
          )
  
      # FF
      elif args.model == "ff_native":
          model = torch.nn.Sequential(
              torch.nn.Linear(
                  args.input_width, args.hidden_width, bias=False
              ),
              torch.nn.GELU(),
              torch.nn.Linear(
                  args.hidden_width, args.output_width, bias=False
              ),
          )
      elif args.model == "ff_bmm":
          model = FF(
              args.input_width, args.hidden_width, args.output_width
          )
      elif args.model == "ff_sparse":
          raise NotImplementedError("Sparse FF is not implemented yet.")
  
      # FFF
      elif args.model == "fff_native":
          raise NotImplementedError(
              "Native FFF cannot be implemented yet -- could YOU be the hero of the PyTorch folk?"
          )
      elif args.model == "fff_bmm":
          model = FFF(args.input_width, args.depth, args.output_width)
      elif args.model == "fff_sparse":
          model = FFF_Sparse(
              args.input_width, args.output_width, args.depth
          )
  
      if not cpu:
          return model.to(torch.float32).cuda()
      else:
          return model.to(torch.float32)
  
file_path: bert-ultrafast/benchmark_pytorch/main.py
imports: import torch
  import warnings
  import sqlite3
  from moe import MoENetwork
  from fff import FF
  from fff import FFF
  from fff import FFF_Sparse
  import time
  import argparse
  import torch._dynamo
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: init_model() and returns: model, model