Mixed precision training leads to NaN-loss

I've been following the Mixed Precision Guide. Hence, I am setting:

keras.mixed_precision.set_global_policy(mixed_precision)

and wrap the optimizer like so:

if mixed_precision.startswith('mixed'):
    logger.info(f'Using LossScaleOptimizer for mixed-precision policy "{mixed_precision}"')
    optimizer = keras.mixed_precision.LossScaleOptimizer(optimizer)

My model has simple Dense layer as an output which I set to 'float32`

# Set dtype explicitly in last layer for mixed-precision training (float32 for numeric stability).
self.output_dense = layers.Dense(vocab_size, dtype=tf.float32)

and a custom train_step() implementation which I modified to this:

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

However, after a while my loss is still becoming NaN:

Outside I am confirming whether or not the policy has been recognized by the model:

logger.info(f'Mixed-precision policy: {mixed_precision}')
logger.info(f'Compute dtype:          {model.compute_dtype}')
logger.info(f'Variable dtype:         {model.variable_dtype}')

keras.py:216] Mixed-precision policy: mixed_float16
keras.py:217] Compute dtype:          float16
keras.py:218] Variable dtype:         float32

But I can tell that this is the case due to the NaN loss..

Is there anything obvious I am doing wrong or missing out? Any idea how I could track down the issue here?