Downloading and preparing dataset 1.15 GiB (download: 1.15 GiB, generated: 310.08 MiB, total: 1.45 GiB) to /root/tensorflow_datasets/mslr_web/10k_fold1/1.2.0...Dl Completed...: 100%1/1 [01:55<00:00, 46.88s/ url]Dl Size...: 100%1176/1176 [01:55<00:00, 26.24 MiB/s]Extraction completed...: 100%15/15 [01:55<00:00, 5.52s/ file]Dataset mslr_web downloaded and prepared to /root/tensorflow_datasets/mslr_web/10k_fold1/1.2.0. Subsequent calls will reuse this data.Epoch 1/3/usr/local/lib/python3.10/dist-packages/keras/src/engine/functional.py:642: UserWarning: Input dict contained keys ['_mask'] which did not match any model input. They will be ignored by the model.  inputs = self._flatten_to_reference_inputs(inputs)---------------------------------------------------------------------------ValueError                                Traceback (most recent call last)[<ipython-input-3-d669266a645f>](https://localhost:8080/#) in <cell line: 41>()     39     loss=tfr.keras.losses.SoftmaxLoss(),     40     metrics=tfr.keras.metrics.get("ndcg", topn=5, name="NDCG@5"))---> 41 model.fit(ds, epochs=3)1 frames[/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py](https://localhost:8080/#) in tf__train_function(iterator)     13                 try:     14                     do_return = True---> 15                     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)     16                 except:     17                     do_return = FalseValueError: in user code:    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1377, in train_function  *        return step_function(self, iterator)    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1360, in step_function  **        outputs = model.distribute_strategy.run(run_step, args=(data,))    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1349, in run_step  **        outputs = model.train_step(data)    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1126, in train_step        y_pred = self(x, training=True)    File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler        raise e.with_traceback(filtered_tb) from None    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/input_spec.py", line 280, in assert_input_compatibility        raise ValueError(    ValueError: Exception encountered when calling layer 'model' (type Functional).        Input 0 of layer "batch_normalization_1" is incompatible with the layer: expected axis 2 of input shape to have value 1, but received input with shape (None, None, 136)        Call arguments received by layer 'model' (type Functional):      • inputs={'_mask': 'tf.Tensor(shape=(None, None), dtype=bool)', 'doc_id': 'tf.Tensor(shape=(None, None), dtype=int64)', 'float_features': 'tf.Tensor(shape=(None, None, 136), dtype=float64)', 'query_id': 'tf.Tensor(shape=(None,), dtype=string)'}      • training=True      • mask=None

tensorflow==2.14.0tensorflow-datasets==4.9.3tensorflow-ranking==0.5.3tfds-nightly==4.9.3.dev202310060044