def train_step(self, data):
    (targets, _) = self.compute_targets(data)
    with tf.GradientTape() as tape:
        logits = self([data, targets], training=True)
        if self._hparams['decoding_task']:
            avg_loss = self._caption_loss(targets, logits)
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(avg_loss, trainable_vars)
        (gradients, global_norm) = tf.clip_by_global_norm(gradients, self._hparams['clip_norm'])
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    self.model_metrics['loss'].update_state(avg_loss)
    self.model_metrics['global_norm'].update_state(global_norm)
    train_metrics = ['loss', 'global_norm']
    return {m: self.model_metrics[m].result() for m in train_metrics}
def train_step(self, data):
    (targets, _) = self.compute_targets(data)
    with tf.GradientTape() as tape:
        logits = self([data, targets], training=True)
        if self._hparams['decoding_task']:
            avg_loss = self._caption_loss(targets, logits)
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(avg_loss, trainable_vars)
        (gradients, global_norm) = tf.clip_by_global_norm(gradients, self._hparams['clip_norm'])
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    self.model_metrics['loss'].update_state(avg_loss)
    self.model_metrics['global_norm'].update_state(global_norm)
    train_metrics = ['loss', 'global_norm']
    return {m: self.model_metrics[m].result() for m in train_metrics}