code: def get_lm_model(args, device, config):
      """Get language model(based on GPT-2) used for sequence prediction."""
  
      ninp = config["ninp"]
      nhead = config["nhead"]
      initrange = config["initrange"]
      dropout = config["dropout"]
      vocab_size = config["vocab_size"]
      nhid = config["nhid"]
      ndecoder = config["num_decoder_layers"]
      is_moe = config.get("is_moe", False)
      num_local_experts = config.get("num_local_experts", 1)
  
      if args.lazy_construction:
          layers = [
              LazyModule(
                  lambda: transformer_lm.EmbeddingLayer(
                      vocab_size, ninp, initrange
                  )
              ),
              LazyModule(
                  lambda: transformer_lm.PositionalEncodingLayer(
                      ninp, dropout
                  )
              ),
          ]
          for _ in range(ndecoder):
              layers.append(
                  LazyModule(
                      lambda: transformer_lm.TransformerDecoderLayer(
                          ninp,
                          nhead,
                          nhid,
                          dropout,
                          is_moe,
                          num_local_experts,
                      )
                  )
              )
  
          layers.append(
              LazyModule(
                  lambda: transformer_lm.LinearLayer(
                      ninp, vocab_size, initrange
                  )
              )
          )
          model = layers
      else:
          model = transformer_lm.TransformerLM(
              vocab_size,
              ninp,
              nhead,
              nhid,
              dropout,
              initrange,
              ndecoder,
              is_moe,
              num_local_experts,
          ).to(device)
  
      return model
  
file_path: fairscale/benchmarks/utils.py
imports: import argparse
  from functools import reduce
  import logging
  import operator
  import datasets.wikitext2_data
  from models import transformer_lm
  import numpy
  import torch
  from torch.optim import Adam
label_desc: Use image node_id 7 for calling a global function with example usage: get_lm_model(args, device, config) and returns: model