i want to get the same number classes in output

[othman8G3]

when i trained my model on new dataset (number classes :18) with transfer learning, but i get in output 406 classes or 181 number classes. so how can i get always just 18 number classes in OUTPUT.

CODE :

def build_model(self, n_categories):
   #n_categories = 18
    input_word_ids = tf.keras.Input(shape=(self.MAX_LEN,), dtype=tf.int32, name='input_word_ids')
    input_mask = tf.keras.Input(shape=(self.MAX_LEN,), dtype=tf.int32, name='input_mask')
    input_type_ids = tf.keras.Input(shape=(self.MAX_LEN,), dtype=tf.int32, name='input_type_ids')

    roberta_model = TFRobertaModel.from_pretrained(self.MODEL_NAME, num_labels = 18, output_hidden_states = False, output_hidden_states = False)

    x = roberta_model(input_word_ids, attention_mask=input_mask, token_type_ids=input_type_ids)
    x = x[0]

    x = tf.keras.layers.Dropout(0.1)(x)
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(256, activation='relu')(x)
    x = tf.keras.layers.Dense(n_categories, activation='softmax')(x)

    model = tf.keras.Model(inputs=[input_word_ids, input_mask, input_type_ids], outputs=x)
    model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

    return model