code: def create_model(vocab_size, args):
      transformer_config = LSTransformer.get_config(
          model="transformer-base",
          max_batch_tokens=2048,
          max_seq_len=512,
          vocab_size=vocab_size,
          padding_idx=0,
          num_encoder_layer=6,
          num_decoder_layer=6,
          fp16=True,
          local_rank=0,
      )
      model = LSTransformer(transformer_config)
      if args.enable_quant:
          model.apply(enable_quant)
      else:
          model.apply(disable_quant)
      model.to(dtype=torch.half, device=torch.device("cuda:0"))
      return model
  
file_path: lightseq/examples/training/custom/run.py
imports: import argparse
  import torch
  from transformers import BertTokenizer
  from lightseq.training import LSTransformer, LSCrossEntropyLayer, LSAdam
  from lightseq.training.ops.pytorch.quantization import disable_quant, enable_quant
label_desc: Use image node_id 2 for calling a global function with example usage: create_model(vocab_size, args) and returns: model