Train step with gradient noise addition. Link to blog: https://medium.com/@afagarap/avoiding-the-vanishing-gradients-problem-96183fd03343?source=friends_link&sk=45cee1239b479c61373e8ff139f89bcc

gradient-noise-addition.py

def train_step(model, loss, features, labels, epoch):
  with tf.GradientTape() as tape:
    logits = model(features)
    train_loss = loss(logits, labels)
  gradients = tape.gradient(train_loss, model.trainable_variables)
  stddev = 1 / ((1 + epoch)**0.55)
  gradients = [tf.add(gradient, tf.random.normal(stddev=stddev, mean=0., shape=gradient.shape)) for gradient in gradients]
  model.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return train_loss, gradients