Unexpected behavior when subclass Tensorflow Keras Model with 3D input for dense layer

Question

I am comparing two simple single dense layer regressors. The two regressors are only different in a way that the 2nd one I pass in index of data in the call function rather than real data and use index to retrieve data for training.

class DenseRegressorV1(tf.keras.Model):
    def __init__(
        self,
        *args,
        **kwargs,
    ):
        super(DenseRegressorV1, self).__init__(*args, **kwargs)
        self.dense_layer    = layers.Dense(units=1, name="dense")

    def call(self, inputs):
        logits = self.dense_layer(inputs)
        return logits
    
class DenseRegressorV2(tf.keras.Model):
    def __init__(
        self,
        data,
        *args,
        **kwargs,
    ):

        super(DenseRegressorV2, self).__init__(*args, **kwargs)
        self.data           = data
        self.dense_layer    = layers.Dense(units=1, name="dense")

    def call(self, inputs_idx):
        inputs = tf.gather(self.data, inputs_idx)
        logits = self.dense_layer(inputs)
        return logits

The input data is in 3D. The y is a simple linear sum of x with some noise which is easy to predict. The data is generated as below:

x_train_3D = tf.random.uniform((80000, 80, 20))
y_train_3D = tf.expand_dims(tf.reduce_sum(x_train_3D, axis=2), axis=2) + 2 * tf.random.uniform((80000, 80, 1))
x_train_3D_idx = np.array(range(0, len(x_train_3D)))

The unexpected behavior is that the V1 seems doing correct training and the val_loss converges to a very small number but V2 does not, although I think they should perform the same.

I use below code for training:

def train_model(model, x_train, y_train, num_epochs, batch_size, learning_rate):
    optimizer = keras.optimizers.Adam(learning_rate)
    tf.keras.metrics.RootMeanSquaredError()
    model.compile(
        optimizer=optimizer,
        loss=tf.keras.losses.mse,
        metrics=['mse'],
    )
    early_stopping = keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=20, restore_best_weights=True
    )
    model_history = model.fit(
        x=x_train,
        y=y_train,
        epochs=num_epochs,
        batch_size=batch_size,
        validation_split=0.18,
        callbacks=[early_stopping]
    )
    return model_history


dr_model_v1   = DenseRegressorV1()
dr_model_v2   = DenseRegressorV2(x_train_3D)

train_model(dr_model_v1, x_train_3D, y_train_3D, num_epochs=500, batch_size=1000, learning_rate=0.01) 
train_model(dr_model_v2, x_train_3D_idx, y_train_3D, num_epochs=500, batch_size=1000, learning_rate=0.01)

So my questions are:

1. What causes the difference in the training and how we can fix for the V2 model?
2. If I flatten the train data into 2D before feeding into the model, I can get expected results for V2 model. Why the shape is making a difference here?

How I transformed to 2D for training with V2

x_train_2D = tf.reshape(x_train_3D, (x_train_3D.get_shape()[0] * x_train_3D.get_shape()[1], x_train_3D.get_shape()[2]))
y_train_2D = tf.reshape(y_train_3D, (y_train_3D.get_shape()[0] * y_train_3D.get_shape()[1], y_train_3D.get_shape()[2]))