code: def create_criterion():
      ce_config = LSCrossEntropyLayer.get_config(
          max_batch_tokens=2048,
          padding_idx=0,
          epsilon=0.0,
          fp16=True,
          local_rank=0,
      )
      loss_fn = LSCrossEntropyLayer(ce_config)
      loss_fn.to(dtype=torch.half, device=torch.device("cuda:0"))
      return loss_fn
  
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 3 for calling a global function with example usage: create_criterion() and returns: loss_fn