Image Captioning with Tensorflow outputting the same caption for all images

I am developing an Image Captioning model using the Flick8k dataset with TensorFlow and I am having an issue where the model is outputting the same caption for every image.

I am using the Xception network to first extract feature vectors for all the images. Then I am creating a TF dataset with the feature vectors and the cleaned and vectorized captions.

The following is the code to generate the dataset, where encoded_img_train/test is a list of feature vectors and cap_train/test is a list of vectorized captions.

def create_dataset(images, captions, batch_size):
  def map_func(decoder_input, decoder_output, feature):
    x = {'decoder_in': decoder_input, 'encoder_in': feature}
    y = {'decoder_out': decoder_output}
    return x, y
  
  decoder_input = []
  decoder_output = []
  for caption in captions:
    in_seq = [token_start] + caption[:-1]
    out_seq = caption[1:] + [token_end]
    # Add padding to input captions
    in_seq = tf.keras.preprocessing.sequence.pad_sequences([in_seq], maxlen=38, padding='post')[0]
    # Add padding to output captions
    out_seq = tf.keras.preprocessing.sequence.pad_sequences([out_seq], maxlen=38, padding='post')[0]

    decoder_input.append(in_seq)
    decoder_output.append(out_seq)

  dataset = tf.data.Dataset.from_tensor_slices((decoder_input, decoder_output, images))
  dataset = dataset.map(map_func)
  dataset = dataset.repeat()
  dataset = dataset.shuffle(1000).batch(64)
  dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

  return dataset

train_dataset = create_dataset(encoded_img_train,cap_train,BATCH_SIZE)
val_dataset = create_dataset(encoded_img_test,cap_test,BATCH_SIZE)

This is my model

inputs1 = Input(shape=(2048,), name='encoder_in')
fe1 = tf.keras.layers.Dropout(0.5)(inputs1)
fe2 = tf.keras.layers.Dense(units, activation='relu', name='encoder_dense')(inputs1)

inputs2 = tf.keras.layers.Input(shape=(None,), name='decoder_in')
se1 = tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim)(inputs2)
decoder_gru1 = GRU(units, name='decoder_gru1',
                   return_sequences=True)(se1)
decoder_gru2 = GRU(units, name='decoder_gru2',
                   return_sequences=True)(decoder_gru1)
decoder_gru3 = GRU(units, name='decoder_gru3',
                   return_sequences=True)(decoder_gru2)
decoder_dense = Dense(vocab_size, activation='linear', name='decoder_out')(decoder_gru3)

model = tf.keras.models.Model(inputs=[inputs1, inputs2], outputs=decoder_dense)
model.summary()

This is how I train my model

optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)

model.compile(loss='sparse_categorical_crossentropy', optimizer=optimizer)

epochs = 20
history = model.fit(train_dataset,validation_data=val_dataset, epochs=epochs, validation_steps=3, steps_per_epoch=steps_per_epoch)

And finally this is how I generate predictions from images

def generate_caption(image_id,true_caption,max_tokens=30):

    encoder_input = features[image_id]
    encoder_input = np.expand_dims(encoder_input, axis=0)

    shape = (1, max_tokens)
    decoder_input = np.zeros(shape=shape, dtype=np.int)

    token_id = token_start
  
    output=[]

    count_tokens = 0

    while token_id != token_end and count_tokens < max_tokens:
        
        decoder_input[0, count_tokens] = token_id

        input_data ={'encoder_in':encoder_input ,'decoder_in': decoder_input}
        
        predict = model.predict(input_data)
        
        token_id = np.argmax(predict[0, count_tokens, :])
        
        output.append(token_id)
        
        count_tokens += 1
    
    print('Predicted caption',tokenizer.sequences_to_texts([output]))
    print(len(output))
    print('True captions',tokenizer.sequences_to_texts([true_caption]))
    
    img = mpimg.imread(image_path+image_id)
    imgplot = plt.imshow(img)
    plt.show()
    
generate_caption(img_test[1],cap_test[1])

During training the training loss stabilizes at around 2.17 and when generating predictions on images the following is the output:

Is the overall architecture of this solution wrong? or is there something I am missing to make this work well?