def METHOD_NAME(self, data):
    x, y, sample_weight = self.unpack_training_data(data)
    with tf.GradientTape() as tape:
        (potentials, sequence_length, kernel), decoded_sequence, *_ = self(
            x, training=True, return_crf_internal=True
        )
        crf_loss = self.compute_crf_loss(
            potentials, sequence_length, kernel, y, sample_weight
        )
        loss = crf_loss + tf.reduce_sum(self.losses)
    gradients = tape.gradient(loss, self.trainable_variables)
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
    # Update metrics (includes the metric that tracks the loss)
    self.compiled_metrics.update_state(y, decoded_sequence)
    # Return a dict mapping metric names to current value
    orig_results = {m.name: m.result() for m in self.metrics}
    crf_results = {"loss": loss, "crf_loss": crf_loss}
    return {**orig_results, **crf_results}