title: Tensorflow strategy error: Mixing different tf.distribute.Strategy objects with overrides train_step keras model
body: <p>I am having difficulties to implement a distributed strategy on my TensorFlow v2 model and I am looking for some help.<br />
  The main difference with the documentation is that my model overrides keras train_step function.<br />
  I got the distribute strategy to work with simple model or with custom training loop but I would like to avoid the latest as I prefer using model.fit which let me use keras callbacks out of the box.</p>
  <p>There is a link to a collab model with custom train_step similar to my model <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/generative/ipynb/dcgan_overriding_train_step.ipynb" rel="nofollow noreferrer">https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/generative/ipynb/dcgan_overriding_train_step.ipynb</a></p>
  <p>To reproduce the error, I modified the init of the model to get the network initialization and compile steps inside the strategy scope</p>
  
  <pre class="lang-py prettyprint-override"><code>class GAN(keras.Model):
      def __init__(self, latent_dim):
          super(GAN, self).__init__()
          self.latent_dim = latent_dim
  
          strategy = tf.distribute.MirroredStrategy()
          with strategy.scope():
            
            self.discriminator = keras.Sequential(
                [
                    keras.Input(shape=(28, 28, 1)),
                    layers.Conv2D(64, (3, 3), strides=(2, 2), padding=&quot;same&quot;),
                    layers.LeakyReLU(alpha=0.2),
                    layers.Conv2D(128, (3, 3), strides=(2, 2), padding=&quot;same&quot;),
                    layers.LeakyReLU(alpha=0.2),
                    layers.GlobalMaxPooling2D(),
                    layers.Dense(1),
                ],
                name=&quot;discriminator&quot;,
            )
  
            self.generator = keras.Sequential(
                [
                    keras.Input(shape=(self.latent_dim,)),
                    layers.Dense(7 * 7 * 128),
                    layers.LeakyReLU(alpha=0.2),
                    layers.Reshape((7, 7, 128)),
                    layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding=&quot;same&quot;),
                    layers.LeakyReLU(alpha=0.2),
                    layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding=&quot;same&quot;),
                    layers.LeakyReLU(alpha=0.2),
                    layers.Conv2D(1, (7, 7), padding=&quot;same&quot;, activation=&quot;sigmoid&quot;),
                ],
                name=&quot;generator&quot;,
            )
            
            self.compile(
                d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
                g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),
                loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),
            )
  
  
      def compile(self, d_optimizer, g_optimizer, loss_fn):
          super(GAN, self).compile()
          self.d_optimizer = d_optimizer
          self.g_optimizer = g_optimizer
          self.loss_fn = loss_fn
  
      def train_step(self, real_images):
          if isinstance(real_images, tuple):
              print(&quot;debug&quot;)
              real_images = real_images[0]
          # Sample random points in the latent space
          batch_size = tf.shape(real_images)[0]
          random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
  
          # Decode them to fake images
          generated_images = self.generator(random_latent_vectors)
  
          # Combine them with real images
          combined_images = tf.concat([generated_images, real_images], axis=0)
  
          # Assemble labels discriminating real from fake images
          labels = tf.concat(
              [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0
          )
          # Add random noise to the labels - important trick!
          labels += 0.05 * tf.random.uniform(tf.shape(labels))
  
          # Train the discriminator
          with tf.GradientTape() as tape:
              predictions = self.discriminator(combined_images)
              d_loss = self.loss_fn(labels, predictions)
          grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
          self.d_optimizer.apply_gradients(
              zip(grads, self.discriminator.trainable_weights)
          )
  
          # Sample random points in the latent space
          random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
  
          # Assemble labels that say &quot;all real images&quot;
          misleading_labels = tf.zeros((batch_size, 1))
  
          # Train the generator (note that we should *not* update the weights
          # of the discriminator)!
          with tf.GradientTape() as tape:
              predictions = self.discriminator(self.generator(random_latent_vectors))
              g_loss = self.loss_fn(misleading_labels, predictions)
          grads = tape.gradient(g_loss, self.generator.trainable_weights)
          self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))
          return {&quot;d_loss&quot;: d_loss, &quot;g_loss&quot;: g_loss}
  </code></pre>
  <p>Then run the code with:</p>
  <pre class="lang-py prettyprint-override"><code>epochs = 30
  gan = GAN(latent_dim=latent_dim)
  gan.fit(
      dataset, epochs=epochs, callbacks=[GANMonitor(num_img=3, latent_dim=latent_dim)]
  )
  </code></pre>
  <p>The full error message :</p>
  <pre class="lang-py prettyprint-override"><code>RuntimeError: in user code:
  
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:806 train_function  *
          return step_function(self, iterator)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:796 step_function  **
          outputs = model.distribute_strategy.run(run_step, args=(data,))
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:1211 run
          return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2585 call_for_each_replica
          return self._call_for_each_replica(fn, args, kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2945 _call_for_each_replica
          return fn(*args, **kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:789 run_step  **
          outputs = model.train_step(data)
      &lt;ipython-input-12-54a272890fa5&gt;:78 train_step
          zip(grads, self.discriminator.trainable_weights)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:519 apply_gradients
          self._create_all_weights(var_list)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:702 _create_all_weights
          _ = self.iterations
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:709 __getattribute__
          return super(OptimizerV2, self).__getattribute__(name)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:836 iterations
          aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:1040 add_weight
          aggregation=aggregation)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/base.py:750 _add_variable_with_custom_getter
          **kwargs_for_getter)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer_utils.py:145 make_variable
          shape=variable_shape if variable_shape else None)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:260 __call__
          return cls._variable_v1_call(*args, **kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:221 _variable_v1_call
          shape=shape)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:67 getter
          return captured_getter(captured_previous, **kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2024 creator_with_resource_vars
          created = self._create_variable(next_creator, **kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/mirrored_strategy.py:450 _create_variable
          values.MirroredVariable, values.SyncOnReadVariable, **kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_utils.py:291 create_mirrored_variable
          value_list = real_mirrored_creator(**kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/mirrored_strategy.py:443 _real_mirrored_creator
          v = next_creator(**kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/variables.py:67 getter
          return captured_getter(captured_previous, **kwargs)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2856 creator
          _require_strategy_scope_strategy(strategy)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:344 _require_strategy_scope_strategy
          _wrong_strategy_scope(strategy, context)
      /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:322 _wrong_strategy_scope
          (context.strategy, strategy))
  
      RuntimeError: Mixing different tf.distribute.Strategy objects: &lt;tensorflow.python.distribute.mirrored_strategy.MirroredStrategy object at 0x7ff32ba58f60&gt; is not &lt;tensorflow.python.distribute.distribute_lib._DefaultDistributionStrategy object at 0x7ff3c0ca0320&gt;
  </code></pre>