title: Replacing gradient calculation of loss function in tensorflow 2.0
body: <p>I would like to replace to gradient function for a loss function in <code>tensorflow 2.0</code>.</p>
  
  <p>Say for example I have a loss function which looks like:</p>
  
  <pre><code>def loss_function(prediction):
      # do some standard tensorflow things here
      return loss
  </code></pre>
  
  <p>I then apply the gradients using the <code>tf.GradientTape</code> method i.e.</p>
  
  <pre><code>with tf.GradientTape() as tape:
  
       prediction = model(input)
       loss = loss_function(prediction)
  
  gradients = tf.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  </code></pre>
  
  <p>My issue is that I want to change the gradient computation and explicitly calculate it myself for just the <code>loss_function()</code> that's currently automatically computed. </p>
  
  <p>I would guess this has something to do with the <code>@tf.custom_gradient</code> decorator, but unsure how I can make it work for the loss.</p>
  
  <p>I am using a custom training loop as apposed to sequential/functional api.</p>