code: def get_model_loss(ctx, model, pretrained, dataset_name, dtype, ckpt_dir=None, start_step=None):
      
      # model
      model, vocabulary = nlp.model.get_model(model,
                                              dataset_name=dataset_name,
                                              pretrained=pretrained, ctx=ctx)
  
      if not pretrained:
          model.initialize(init=mx.init.Normal(0.02), ctx=ctx)
      model.cast(dtype)
  
      if ckpt_dir and start_step:
          param_path = os.path.join(ckpt_dir, '%07d.params'%start_step)
          model.load_parameters(param_path, ctx=ctx)
          logging.info('Loading step %d checkpoints from %s.', start_step, param_path)
  
      model.hybridize(static_alloc=True)
  
      # losses
      nsp_loss = mx.gluon.loss.SoftmaxCELoss()
      mlm_loss = mx.gluon.loss.SoftmaxCELoss()
      nsp_loss.hybridize(static_alloc=True)
      mlm_loss.hybridize(static_alloc=True)
  
      return model, nsp_loss, mlm_loss, vocabulary
docs: Get model for pre-training.