code: def BERTPretrainingLoss():
      nsp_loss = [
          tl.Select([0, 2, 3], n_in=6),
          tl.WeightedCategoryCrossEntropy(),
      ]
      mlm_loss = [
          tl.Select([1, 4, 5], n_in=6),
          tl.WeightedCategoryCrossEntropy(),
      ]
      return tl.Serial(tl.Branch(nsp_loss, mlm_loss), tl.Add())
  
imports: import os
  import pathlib
  import urllib
  import zipfile
  import gin
  import jax
  import tensorflow
  from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: BERTPretrainingLoss() and returns: tl