class_docstr: Wrapper around the base model for custom training logic.
code: def train_step(self, data):
      x, y, sample_weight = unpack_data(data)
  
      with tf.GradientTape() as tape:
          crf_loss, internal_losses = self.compute_loss(
              x, y, sample_weight, training=True
          )
          total_loss = crf_loss + internal_losses
  
      gradients = tape.gradient(total_loss, self.trainable_variables)
      self.optimizer.apply_gradients(
          zip(gradients, self.trainable_variables)
      )
  
      return {"crf_loss": crf_loss, "internal_losses": internal_losses}
  
imports: import itertools
  import os
  import math
  import tempfile
  import pytest
  import numpy
  import tensorflow
  from tensorflow_addons.layers.crf import CRF
  from tensorflow_addons.text.crf import crf_log_likelihood
  from tensorflow_addons.utils import test_utils
label_desc: Use image node_id 4 for calling the ModelWithCRFLoss obj's underlying member method code with example usage: obj.train_step(data) and returns: dict