Ive recently tried to train a Python TensorFlow LSTM neural network on Google Colab (Since they provide a well TPU service), but was unable to get the model to run on my own PC. Running the same training code on my local machine produces a compatible model, although at a much slower rate.

The main error I keep on getting is when I try to load the model from Google Colab on my machine. It results in the error below:

ValueError: Unrecognized keyword arguments passed to LSTM: {‘time_major’: False}

The loading code:
model = tf.keras.models.load_model('main_model.h5')

Model Trainer:

model = Sequential([
    Embedding(input_dim=vocab_size, output_dim=embedding_dim, mask_zero=True),
    LSTM(units=lstm_units, return_sequences=True),
    TimeDistributed(Dense(3))
])
model.compile(optimizer='adam', loss='mean_squared_error')

*I Also save the model file as .h5 instead of keras since keras returns an unkown error when loading.

I think the problem mostly arises from the version difference between Keras/Tensorflow on google Colab and the local verisons I have. I have updated versions (Keras 3.4.1 and TF 2.17), while google has Keras/TF both as version 2.15. But I may be wrong.

I mostly focused on the version issue, so i’ve tried lowering my local version to both 2.15, but TensorFlow does not support 2.15 pip installitions even though direct Guthub links. Ive also tried updating the versions on google Colab, but when I do the TPU support and speed disappears. Then I tested only exporting checkpints/model weights and compiling the network on my local machine but many version errors occured as well. Finally, I tried manually modifying the model configs to remove the {'time_major': False}, but this caused an unknown model signature error.

Are some other things I can do to fix this?

Btw im relatively new to TensorFlow Neural Network and posting on Stackoverflow, any help is much appreciated.
Thank you.