title: ValueError: Dimensions must be equal, but are 96 and 256 in tpu on tensorflow
body: <p>I am trying to create a mnist gan which will use tpu.
  I copied the gan code from <a href="https://www.tensorflow.org/tutorials/generative/dcgan" rel="nofollow noreferrer">here</a>.</p>
  <p>Then i made some of my own modifications to run the code on tpu.for making changes i followed <a href="https://www.tensorflow.org/guide/tpu" rel="nofollow noreferrer">this tutorial</a> which shows how to us tpu on tensorflow on tensorflow website.</p>
  <p>but thats not working and raising an error here is my code.</p>
  <pre class="lang-py prettyprint-override"><code># -*- coding: utf-8 -*-
  &quot;&quot;&quot;Untitled13.ipynb
  
  Automatically generated by Colaboratory.
  
  Original file is located at
      https://colab.research.google.com/drive/1gbHDaCeFUCGDkkNgAPjGFQIDvZ5NxVfr
  &quot;&quot;&quot;
  
  # Commented out IPython magic to ensure Python compatibility.
  # %tensorflow_version 2.x
  
  import tensorflow as tf
  import numpy as np
  
  resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
  tf.config.experimental_connect_to_cluster(resolver)
  # This is the TPU initialization code that has to be at the beginning.
  tf.tpu.experimental.initialize_tpu_system(resolver)
  print(&quot;All devices: &quot;, tf.config.list_logical_devices('TPU'))
  
  strategy = tf.distribute.TPUStrategy(resolver)
  
  import glob
  import matplotlib.pyplot as plt
  import numpy as np
  import os
  import PIL
  from tensorflow.keras import layers
  import time
  
  from IPython import display
  
  (train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()
  
  train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
  train_images = (train_images - 127.5) / 127.5  # Normalize the images to [-1, 1]
  
  BUFFER_SIZE = 60000
  BATCH_SIZE = 256
  
  # Batch and shuffle the data
  train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
  
  def make_generator_model():
      model = tf.keras.Sequential()
      model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,)))
      model.add(layers.BatchNormalization())
      model.add(layers.LeakyReLU())
  
      model.add(layers.Reshape((7, 7, 256)))
      assert model.output_shape == (None, 7, 7, 256)  # Note: None is the batch size
  
      model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
      assert model.output_shape == (None, 7, 7, 128)
      model.add(layers.BatchNormalization())
      model.add(layers.LeakyReLU())
  
      model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
      assert model.output_shape == (None, 14, 14, 64)
      model.add(layers.BatchNormalization())
      model.add(layers.LeakyReLU())
  
      model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
      assert model.output_shape == (None, 28, 28, 1)
  
      return model
  
  def make_discriminator_model():
      model = tf.keras.Sequential()
      model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same',
                                       input_shape=[28, 28, 1]))
      model.add(layers.LeakyReLU())
      model.add(layers.Dropout(0.3))
  
      model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
      model.add(layers.LeakyReLU())
      model.add(layers.Dropout(0.3))
  
      model.add(layers.Flatten())
      model.add(layers.Dense(1))
  
      return model
  
  # This method returns a helper function to compute cross entropy loss
  cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
  
  EPOCHS = 50
  noise_dim = 100
  num_examples_to_generate = 16
  
  # You will reuse this seed overtime (so it's easier)
  # to visualize progress in the animated GIF)
  seed = tf.random.normal([num_examples_to_generate, noise_dim])
  
  def generate_and_save_images(model, epoch, test_input):
    # Notice `training` is set to False.
    # This is so all layers run in inference mode (batchnorm).
    predictions = model(test_input, training=False)
  
    fig = plt.figure(figsize=(4, 4))
  
    for i in range(predictions.shape[0]):
        plt.subplot(4, 4, i+1)
        plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')
        plt.axis('off')
  
    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
    plt.show()
  
  def train(dataset, epochs):
    for epoch in range(epochs):
      start = time.time()
  
      for image_batch in (dataset):
        strategy.run(train_step, args=(image_batch,))
  
      # Produce images for the GIF as you go
      display.clear_output(wait=True)
      generate_and_save_images(generator,
                               epoch + 1,
                               seed)
  
      # Save the model every 15 epochs
      if (epoch + 1) % 15 == 0:
        checkpoint.save(file_prefix = checkpoint_prefix)
  
      print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))
  
    # Generate after the final epoch
    display.clear_output(wait=True)
    generate_and_save_images(generator,
                             epochs,
                             seed)
  
  def generator_loss(fake_output):
      return cross_entropy(tf.ones_like(fake_output), fake_output)
  
  def discriminator_loss(real_output, fake_output):
      real_loss = cross_entropy(tf.ones_like(real_output), real_output)
      fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
      total_loss = real_loss + fake_loss
      return total_loss
  
  # Notice the use of `tf.function`
  # This annotation causes the function to be &quot;compiled&quot;.
  @tf.function
  def train_step(images):
      noise = tf.random.normal([BATCH_SIZE, noise_dim])
  
      with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)
  
        real_output = discriminator(images, training=True)
        fake_output = discriminator(generated_images, training=True)
  
        fake_output_0 = discriminator(generated_images, training=True)
  
        gen_loss = generator_loss(fake_output_0)
        disc_loss = discriminator_loss(real_output, fake_output)
  
      gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
      gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
  
      generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
      discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
  
  with strategy.scope():
    generator = make_generator_model()
    generator_optimizer = tf.keras.optimizers.Adam(1e-4)
  
    discriminator = make_discriminator_model()
    discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)
  
    checkpoint_dir = './training_checkpoints'
    checkpoint_prefix = os.path.join(checkpoint_dir, &quot;ckpt&quot;)
    checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                    discriminator_optimizer=discriminator_optimizer,
                                    generator=generator,
                                    discriminator=discriminator)
    
    train(train_dataset, EPOCHS)
  </code></pre>
  <p>and the final output is (not showing whole output cause i am in colab and i do not want copy output pf each cell one by one)</p>
  <pre><code>ValueError: Dimensions must be equal, but are 96 and 256 for '{{node add}} = AddV2[T=DT_FLOAT](binary_crossentropy/weighted_loss/Mul, binary_crossentropy_1/weighted_loss/Mul)' with input shapes: [96], [256].
  </code></pre>