code: def new_train_step(
      self,
      loss_function: Callable[[tf.Tensor], tf.Tensor],
      batch_size: int,
      apply_gradients=True,
  ):
      transformer = self.transformer
      optimizer_encoder = self.optimizer_encoder
      optimizer_decoder = self.optimizer_decoder
  
      train_step_signature = [
          tf.TensorSpec(shape=(batch_size, None), dtype=tf.int32),
          tf.TensorSpec(shape=(batch_size, None), dtype=tf.int32),
          tf.TensorSpec(shape=(batch_size, None), dtype=tf.int32),
      ]
  
      @tf.function(input_signature=train_step_signature)
      def train_step(inp, sent_ids, tar):
          tar_inp = tar[:, :-1]
          tar_real = tar[:, 1:]
          (
              enc_padding_mask,
              combined_mask,
              dec_padding_mask,
          ) = create_masks(inp, tar_inp)
          with tf.GradientTape(persistent=True) as tape:
              predictions, _ = transformer(
                  inp,
                  sent_ids,
                  tar_inp,
                  True,
                  enc_padding_mask,
                  combined_mask,
                  dec_padding_mask,
              )
              loss = loss_function(tar_real, predictions)
          if apply_gradients:
              gradients_decoder = tape.gradient(
                  loss, transformer.decoder.trainable_variables
              )
              gradients_encoder = tape.gradient(
                  loss, transformer.encoder.trainable_variables
              )
              optimizer_decoder.apply_gradients(
                  zip(
                      gradients_decoder,
                      transformer.decoder.trainable_variables,
                  )
              )
              optimizer_encoder.apply_gradients(
                  zip(
                      gradients_encoder,
                      transformer.encoder.trainable_variables,
                  )
              )
  
          return loss
  
      return train_step
  
imports: import os
  import pickle
  from typing import Callable
  from typing import Dict, Union
  import numpy
  import tensorflow
  from headliner.model.bert_model import Transformer, create_masks
  from headliner.model.summarizer import Summarizer
  from headliner.preprocessing.bert_vectorizer import BertVectorizer
  from headliner.preprocessing.preprocessor import Preprocessor
  from headliner.utils.logger import get_logger
label_desc: Use image node_id 4 for calling the BertSummarizer obj's underlying member method code with example usage: obj.new_train_step(loss_function, batch_size, apply_gradients) and returns: train_step, loss