model with longformer encoder cannot be saved due to OperatorNotAllowedInGraphError
Author: rdisipio
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.21.2
- Platform: Linux-5.10.102-99.473.amzn2.x86_64-x86_64-with-glibc2.10 (AWS SageMaker)
- Python version: 3.8.12
- Huggingface_hub version: 0.9.1
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): 2.9.1 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Hi,  I created a NER model which makes use of the longformer encoder. I can train it successfully. However, if I try to save the model like this: tf.keras.models.save_model(self.model, model_path) I get an `OperatorNotAllowedInGraphError` error. The full message is below. As far as I can tell, the function `_pad_to_window_size` calculates `padding_len` which in my execution happens to be a tf.Tensor object. The statement `if padding_len > 0` fails to evaluate, since a tensor cannot be compared to a bool directly. This smells like a bug. Besides, the only purpose of this check seems to be to print out a message in the log, so perhaps it's not strictly necessary.

Cheers,
Riccardo


```
---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
/tmp/ipykernel_29348/3622010497.py in <cell line: 2>()
      1 output_path = "trained_models/jd_parser_baseline_all-longformer"
----> 2 t.save_model(output_path)

~/SageMaker/jd-parser/jd_parser/trainer.py in save_model(self, model_path)
    381         if model_path is None:
    382             model_path = f"trained_models/{self.model.name}"
--> 383         tf.keras.models.save_model(self.model, model_path)
    384 
    385     def upload_to_s3(self, model_local_dir=None):

~/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/contextlib.py in __exit__(self, type, value, traceback)
    118         if type is None:
    119             try:
--> 120                 next(self.gen)
    121             except StopIteration:
    122                 return False

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/transformers/modeling_tf_utils.py in run_call_with_unpacked_inputs(self, *args, **kwargs)
    411 
    412         unpacked_inputs = input_processing(func, config, **fn_args_and_kwargs)
--> 413         return func(self, **unpacked_inputs)
    414 
    415     # Keras enforces the first layer argument to be passed, and checks it through `inspect.getfullargspec()`. This

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/transformers/models/longformer/modeling_tf_longformer.py in call(self, input_ids, attention_mask, head_mask, global_attention_mask, token_type_ids, position_ids, inputs_embeds, output_attentions, output_hidden_states, return_dict, training)
   1728             position_ids,
   1729             inputs_embeds,
-> 1730         ) = self._pad_to_window_size(
   1731             input_ids=input_ids,
   1732             attention_mask=attention_mask,

~/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/transformers/models/longformer/modeling_tf_longformer.py in _pad_to_window_size(self, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds, pad_token_id)
   1814         padding_len = (attention_window - seq_len % attention_window) % attention_window
   1815 
-> 1816         if padding_len > 0:
   1817             logger.info(
   1818                 f"Input ids are automatically padded from {seq_len} to {seq_len + padding_len} to be a multiple of "

OperatorNotAllowedInGraphError: Exception encountered when calling layer "longformer" (type TFLongformerMainLayer).

Using a symbolic `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Call arguments received by layer "longformer" (type TFLongformerMainLayer):
  • args=({'input_ids': 'tf.Tensor(shape=(None, None), dtype=int32)', 'attention_mask': 'tf.Tensor(shape=(None, None), dtype=int32)'},)
  • kwargs={'training': 'False'}
```

### Expected behavior

The model should be saved to disk correctly as it happens with other encoder models.

Comments:
@rdisipio  I tried to reproduce this error, But I was able to save model. I followed [this](https://huggingface.co/docs/transformers/model_doc/longformer#transformers.TFLongformerForTokenClassification.call.example) example to build a simple model.  Can you add some steps to reproduce this bug ? 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.