title: Mixed precision training leads to NaN-loss
body: <p>I've been following the <a href="https://www.tensorflow.org/guide/mixed_precision" rel="nofollow noreferrer">Mixed Precision Guide</a>. Hence, I am setting:</p>
  <pre class="lang-py prettyprint-override"><code>keras.mixed_precision.set_global_policy(mixed_precision)
  </code></pre>
  <p>and wrap the optimizer like so:</p>
  <pre class="lang-py prettyprint-override"><code>if mixed_precision.startswith('mixed'):
      logger.info(f'Using LossScaleOptimizer for mixed-precision policy &quot;{mixed_precision}&quot;')
      optimizer = keras.mixed_precision.LossScaleOptimizer(optimizer)
  </code></pre>
  <p>My model has simple <code>Dense</code> layer as an output which I set to 'float32`</p>
  <pre class="lang-py prettyprint-override"><code># Set dtype explicitly in last layer for mixed-precision training (float32 for numeric stability).
  self.output_dense = layers.Dense(vocab_size, dtype=tf.float32)
  </code></pre>
  <p>and a custom <code>train_step()</code> implementation which I modified to this:</p>
  <pre class="lang-py prettyprint-override"><code>with tf.GradientTape() as tape:
      model_loss = self.loss_fn(
          inputs,
          y_true=y_true,
          mask=mask
      )
  
      is_mixed_precision = isinstance(self.optimizer, mixed_precision.LossScaleOptimizer)
  
      # We always want to return the unmodified model_loss for Tensorboard
      if is_mixed_precision:
          loss = self.optimizer.get_scaled_loss(model_loss)
      else:
          loss = model_loss
  
      gradients = tape.gradient(loss, self.trainable_variables)
  
      if is_mixed_precision:
          gradients = self.optimizer.get_unscaled_gradients(gradients)
  
      return model_loss, gradients
  </code></pre>
  <p>However, after a while my loss is still becoming <code>NaN</code>:</p>
  <p><a href="https://i.stack.imgur.com/RB1hg.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/RB1hg.png" alt="enter image description here" /></a></p>
  <p>Outside I am confirming whether or not the policy has been recognized by the model:</p>
  <pre class="lang-py prettyprint-override"><code>logger.info(f'Mixed-precision policy: {mixed_precision}')
  logger.info(f'Compute dtype:          {model.compute_dtype}')
  logger.info(f'Variable dtype:         {model.variable_dtype}')
  </code></pre>
  <pre class="lang-none prettyprint-override"><code>keras.py:216] Mixed-precision policy: mixed_float16
  keras.py:217] Compute dtype:          float16
  keras.py:218] Variable dtype:         float32
  </code></pre>
  <p>But I can tell that this is the case due to the <code>NaN</code> loss..</p>
  <p>Is there anything obvious I am doing wrong or missing out? Any idea how I could track down the issue here?</p>