title: Keras model doesn't train when using GradientTape
body: <p>I have a class with a model inside:</p>
  
  <pre><code>class MyNetwork:
  
      def __init__(self):
          # layer initializations
          self.optimizer = tf.train.RMSPropOptimizer(0.001)
          self.loss = tf.losses.sigmoid_cross_entropy
  
      def build(self):
          # layer connections
  
          self.model = keras.Model(inputs=[inputs], outputs=[outputs]
          return self.model
  
      @tf.function
      def train_step(self, images, labels):
          with tf.GradientTape() as tape:
              predictions = self.model(images)
              loss = self.loss(labels, predictions)
  
          gradients = tape.gradient(loss, self.model.trainable_variables)
          self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
          return loss, predictions
  </code></pre>
  
  <p>I use the following to build the model:</p>
  
  <pre><code>network = MyNetwork()
  model = network.build()
  </code></pre>
  
  <p>When training with the following lines:</p>
  
  <pre><code>model.compile(tf.train.RMSPropOptimizer(0.001), loss=tf.losses.sigmoid_cross_entropy, metrics=['accuracy']
  model.fit(X, y, epochs=10)
  </code></pre>
  
  <p>The model trains without any issues.</p>
  
  <p>But in a separate run, using the following code:</p>
  
  <pre><code>for i in range(10):
      print("Epoch ", (i))
  
      loss, pred = D.train_step(X, y)
      print(loss)
  </code></pre>
  
  <p>The loss gets stuck after a a few epochs. The calculated accuracy also stops at 0.5.</p>
  
  <p>Does anyone know how a Keras model can be trained using <code>GradientTape</code>?</p>