tensorflow - ステートフル LSTM Tensorflow 無効な Input_h 形状エラー

Question

TensorFlow を使用して、時系列回帰問題でステートフル LSTM を試しています。データセットを共有できないことをお詫び申し上げます。以下は私のコードです。

train_feature = train_feature.reshape((train_feature.shape[0], 1, train_feature.shape[1]))
val_feature = val_feature.reshape((val_feature.shape[0], 1, val_feature.shape[1]))

batch_size = 64

model = tf.keras.Sequential()
model.add(tf.keras.layers.LSTM(50, batch_input_shape=(batch_size, train_feature.shape[1], train_feature.shape[2]), stateful=True))
model.add(tf.keras.layers.Dense(1))

model.compile(optimizer='adam',
              loss='mse',
              metrics=[tf.keras.metrics.RootMeanSquaredError()])

model.fit(train_feature, train_label, 
          epochs=10,
          batch_size=batch_size)

上記のコードを実行すると、最初のエポックの終了後に次のようなエラーが発生します。

InvalidArgumentError:  [_Derived_]  Invalid input_h shape: [1,64,50] [1,49,50]
     [[{{node CudnnRNN}}]]
     [[sequential_1/lstm_1/StatefulPartitionedCall]] [Op:__inference_train_function_1152847]

Function call stack:
train_function -> train_function -> train_function

ただし、batch_size を 1に変更し、モデル トレーニングのコードを次のように変更すると、モデルは正常にトレーニングされます。

total_epochs = 10

for i in range(total_epochs):
    model.fit(train_feature, train_label, 
              epochs=1,
              validation_data=(val_feature, val_label),
              batch_size=batch_size,
              shuffle=False)

    model.reset_states()

ただし、非常に大きなデータ (100 万行) では、batch_size が 1 であるため、モデルのトレーニングに非常に長い時間がかかります。

それで、無効なinput_h形状エラーを取得せずに、バッチサイズが1より大きい（たとえば64）ステートフルLSTMをトレーニングするにはどうすればよいでしょうか？

回答ありがとうございます。