title: Shape of output change during TPU training Keras funcional API
body: <p><em>ValueError: Mismatch between expected batch size and model output batch size. Output shape = (16, 6248), expected output shape = shape (1, 6248)</em></p>
  
  <p>I defined in the last dense layers of my model is (None, 6248)</p>
  
  <p>I can't understandable what is the number 16 in this output. All of code without any number 16.</p>
  
  <p>Desciption:
  i'm use TPU V3 in kaggle for training. <strong>but if i using GPU then i'm not get any Error when i make prediction "model.predict(...)"</strong>. </p>
  
  <p>I create model with Keras functional API</p>
  
  <p>Please make me clearly. Thank!
  Here is my model.</p>
  
  <pre class="lang-py prettyprint-override"><code>    embedding_layer = Embedding(len(word_index) + 1,
                                  EMBEDDING_DIM,
                                  weights=[embedding_matrix],
                                  input_length= max_sequence_len - 1,
                                  trainable=False)
  
      sequence_1_input = Input(shape=(max_sequence_len - 1,), dtype='int32')
      embedded_sequences_1 = embedding_layer(sequence_1_input)
  
      activations = Bidirectional(LSTM(num_lstm, dropout=rate_drop_lstm, recurrent_dropout=rate_drop_lstm, return_sequences=True))(embedded_sequences_1)
      activations = LSTM(num_lstm, dropout=rate_drop_lstm, recurrent_dropout=rate_drop_lstm)(activations)
  
      merged = Dense(num_dense, activation=act)(activations)
      merged = Dropout(rate_drop_dense)(merged)
      merged = BatchNormalization()(merged)
  
      preds = Dense(vocab_size, activation='softmax')(merged)
  
  
      model = Model(inputs=sequence_1_input, outputs=preds)
      model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['categorical_accuracy'])
  </code></pre>
  
  <p>Here is my data:</p>
  
  <pre class="lang-py prettyprint-override"><code>train = predictors[0:10240]
  test = label[0:10240]
  
  dataset = tf.data.Dataset.from_tensor_slices((train, test))
  dataset = dataset.shuffle(200).batch(1024).repeat()
  dataset = dataset.prefetch(AUTO)
  
  #Here is the shape of dataset:
  
  &lt;PrefetchDataset shapes: ((None, 87), (None, 6248)), types: (tf.int32, tf.float32)&gt;
  </code></pre>
  
  <p>Fit model:</p>
  
  <pre class="lang-py prettyprint-override"><code>model.fit(dataset, epochs=100, steps_per_epoch=10240//1024, verbose=1)
  </code></pre>
  
  <p>Summary</p>
  
  <p><a href="https://i.stack.imgur.com/olY8P.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/olY8P.png" alt="enter image description here"></a></p>