Transformer Model isn't fitting well on Time Series Data - tensorflow

Question

I have created two models LSTM, LSTM with Self-Attention. Now I am working on to create my first transformer model. I created it for multivariate time series predictions (many-to-one classification model).
I have the hourly varying data i.e., 8 different features (hour, month, temperature, humidity, windspeed, solar radiations concentration etc.) and with them I am trying to predict the time sequence (energy consumption of a building. So my input has the shape X.shape = (8783, 168, 8) i.e., 8783 time sequences, each sequence contains 168 hourly entries/vectors and each vector contains 8 features. My output has the shape Y.shape = (8783,1) i.e., 8783 sequences each containing 1 output value (i.e., building energy consumption value after every hour).
I took as a model an example from the official keras site. It is created for classification problems, I converted my output to classes n_classes = len(np.unique(Y_train)) = 156
Input shape (X_train) = (8783, 168, 8) Output shape (Y_train) = (8783,1) n_classes = 156
In the softmax activation layer, I set the output to n_classes but it's not fitting well. Below I attach the model and I would like to know:
I)- Have I done something wrong in the model? Is the model architecture is fine? Are there maybe some other parts of the code I need to change for it to work for my problem?
II)- Also, can a transformer at all work on multivariate problems of my kind (8 features input, 1 feature output) or do transformers only work on univariate problems?
def build_transformer_model(input_shape, head_size, num_heads, ff_dim, num_transformer_blocks,
mlp_units, dropout=0, mlp_dropout=0):
inputs = keras.Input(shape=input_shape)
x = inputs
for _ in range(num_transformer_blocks):
# Normalization and Attention
x = layers.LayerNormalization(epsilon=1e-6)(x)
x = layers.MultiHeadAttention(
key_dim=head_size, num_heads=num_heads, dropout=dropout
)(x, x)
x = layers.Dropout(dropout)(x)
res = x + inputs
# Feed Forward Part
x = layers.LayerNormalization(epsilon=1e-6)(res)
x = layers.Conv1D(filters=ff_dim, kernel_size=1, activation="relu")(x)
x = layers.Dropout(dropout)(x)
x = layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
x = x + res
x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
for dim in mlp_units:
x = layers.Dense(dim, activation="relu")(x)
x = layers.Dropout(mlp_dropout)(x)
x = layers.Dense(n_classes, activation="softmax")(x)
return keras.Model(inputs, x)
model_tr = build_transformer_model(input_shape=(X_train.shape[1], X_train.shape[2]), head_size=256,
num_heads=4, ff_dim=4, num_transformer_blocks=4, mlp_units=[128], mlp_dropout=0.4, dropout=0.2)
model_tr.compile(loss="sparse_categorical_crossentropy",
optimizer=keras.optimizers.Adam(learning_rate=0.0001),
metrics=["sparse_categorical_accuracy"],
)
plot_model(model_tr, to_file='model_plot.png', show_shapes=True, show_layer_names=True)
model_tr.summary()
m_tr_history = model_tr.fit(x=X_train, y=Y_train, validation_split=0.15, batch_size=64, epochs=100, verbose = 1)
model_tr.save('halka_transformer.h5')
No. of epochs
Thanks in advance for this valuable assistance.