KeyError                                  Traceback (most recent call last)
[<ipython-input-49-089aabb58b9b>](https://localhost:8080/#) in <cell line: 1>()
----> 1 history = model.fit(x=tf_train_set, validation_data=tf_validation_set, epochs=num_epochs, callbacks=callbacks)

1 frames
[/usr/local/lib/python3.10/dist-packages/transformers/keras_callbacks.py](https://localhost:8080/#) in on_epoch_end(self, epoch, logs)
    217             if self.predict_with_generate:
    218                 if isinstance(batch, dict):
--> 219                     generation_inputs = batch[main_input_name]
    220                     attention_mask = batch.get("attention_mask", None)
    221                 else:

KeyError: None

#### in tf_keras_callback (in func `on_epoch_end`, ~line 191)
main_input_name = None
if self.predict_with_generate:
    # This dense conditional recognizes the case where we have an encoder-decoder model, but
    # avoids getting tangled up when we just have a model with a layer called 'encoder'
    if hasattr(self.model, "encoder") and hasattr(self.model.encoder, "main_input_name"):
        # >>>>>>> If this condition is not satisfied(which is the case currently), `main_input_name` remains None <<<<<<<<
        if self.model.encoder.main_input_name != self.model.main_input_name:
            main_input_name = self.model.encoder.main_input_name
    else:
        main_input_name = getattr(self.model, "main_input_name", "input_ids")

    if self.use_xla_generation and self.generation_function is None:

        def generation_function(inputs, attention_mask):
            return self.model.generate(inputs, attention_mask=attention_mask, **self.generate_kwargs)

        self.generation_function = tf.function(generation_function, jit_compile=True)

prediction_list = []
label_list = []

# The whole predict/generate loop is handled inside this method
for batch in self.eval_dataset:
    if isinstance(batch, tuple):
        batch, labels = batch
    else:
        labels = None
    if self.predict_with_generate:
        if isinstance(batch, dict):
            generation_inputs = batch[main_input_name]  # >>>>>>>>>>>> `main_input_name` remains None here (~line 219) <<<<<<<<<<<<
            attention_mask = batch.get("attention_mask", None)
        else:
            generation_inputs = batch
            attention_mask = None
        if self.use_xla_generation:
            predictions = self.generation_function(generation_inputs, attention_mask=attention_mask)
        else:
            predictions = self.model.generate(
                generation_inputs, attention_mask=attention_mask, **self.generate_kwargs
            )
    

if hasattr(self.model, "encoder") and hasattr(self.model.encoder, "main_input_name"):
        if self.model.encoder.main_input_name != self.model.main_input_name:
            main_input_name = self.model.encoder.main_input_name

if hasattr(self.model, "encoder") and hasattr(self.model.encoder, "main_input_name"):
        main_input_name = self.model.encoder.main_input_name