Exception encountered when calling layer "lstm_cell" (type LSTMCell)

[MNI-S]

I am trying to implement a deep learning model combining CNN+ LSTM in AWS Sagemaker. While trying to add the LSTM layer after CNN, I get the following value error. I also run the same code in Google collab or my Local Machine, but that problem doesn't appear.

---

ValueError Traceback (most recent call last)
/tmp/ipykernel_4202/2122205439.py in <cell line: 3>()
1 tf.keras.backend.clear_session()
2
----> 3 model_06 = get_model()

/tmp/ipykernel_4202/3426630573.py in get_model()
29 x = Reshape((64, 256))(x)
30 print(f"After Reshaping(): {x.shape}")
---> 31 x = (LSTM(512, activation = "relu"))(x)
32
33 # Fully Connected layer

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/keras/layers/rnn/base_rnn.py in call(self, inputs, initial_state, constants, **kwargs)
513
514 if initial_state is None and constants is None:
--> 515 return super(RNN, self).call(inputs, **kwargs)
516
517 # If any of initial_state or constants are specified and are Keras

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
65 except Exception as e: # pylint: disable=broad-except
66 filtered_tb = _process_traceback_frames(e.traceback)
---> 67 raise e.with_traceback(filtered_tb) from None
68 finally:
69 del filtered_tb

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/keras/backend.py in bias_add(x, bias, data_format)
6362 if len(bias_shape) == 1:
6363 if data_format == 'channels_first':
-> 6364 return tf.nn.bias_add(x, bias, data_format='NCHW')
6365 return tf.nn.bias_add(x, bias, data_format='NHWC')
6366 if ndim(x) in (3, 4, 5):

ValueError: Exception encountered when calling layer "lstm_cell" (type LSTMCell).

Shape must be at least rank 3 but is rank 2 for '{{node lstm/lstm_cell/BiasAdd}} = BiasAdd[T=DT_FLOAT, data_format="NCHW"](lstm/lstm_cell/add, lstm/lstm_cell/BiasAdd/ReadVariableOp)' with input shapes: [?,2048], [2048].

Call arguments received by layer "lstm_cell" (type LSTMCell):
• inputs=tf.Tensor(shape=(None, 256), dtype=float32)
• states=('tf.Tensor(shape=(None, 512), dtype=float32)', 'tf.Tensor(shape=(None, 512), dtype=float32)')
• training=None

Here is my complete code:

`
def get_model():

# input layer
inputs = tf.keras.layers.Input(shape=(3, 256, 256), name="input_layer")

# 2nd convolution block
x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', padding='same')(inputs)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)

# 2nd convolution block
x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu', padding='same')(x)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)

# 3rd convolution block
x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation='relu', padding='same')(x)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)

# 4th convolution block
x = tf.keras.layers.Conv2D(filters=128, kernel_size=(3, 3), activation='relu', padding='same')(x)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)

# 5th convolution block
x = tf.keras.layers.Conv2D(filters=256, kernel_size=(3, 3), activation='relu', padding='same')(x)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)

print(f"Before Reshaping(): {x.shape}")

# LSTM
x = Reshape((64, 256))(x)
print(f"After Reshaping(): {x.shape}")
x = (LSTM(512, activation = "relu"))(x)

# Fully Connected layer
x = Flatten()(x)
x = tf.keras.layers.Dense(1024, activation = 'relu')(x)

# Output layer
outputs = tf.keras.layers.Dense(1, activation="sigmoid", name="output_layer")(x)

# Combine the inputs with the outputs into a model
model = tf.keras.Model(inputs, outputs)

return model

`