Cannot save TFSwinForImageClassification as SavedModel
Author: ahmedlone127
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.20.1
- Platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.13
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.0+cu113 (True)
- Tensorflow version (GPU?): 2.8.2 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?:No

### Who can help?

@Rocketknight1 @sgugger



### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction
```
from transformers import AutoFeatureExtractor, TFSwinForImageClassification
from PIL import Image
import requests

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

feature_extractor = AutoFeatureExtractor.from_pretrained(swinModel)
model = TFSwinForImageClassification.from_pretrained(swinModel)

inputs = feature_extractor(images=image, return_tensors="tf")
outputs = model(inputs.pixel_values)
logits = outputs.logits
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = tf.math.argmax(logits,-1).numpy()[0]
print("Predicted class:", model.config.id2label[predicted_class_idx])

class MySwin(TFSwinForImageClassification):
    @tf.function(
        input_signature=[
            {
                "pixel_values": tf.TensorSpec((None, None,None,None), tf.float32, name="serving1_pixel_values"),
            }
        ]
    )
    def serving1(self, inputs):
        outputs = self.call(pixel_values=inputs["pixel_values"])
        return self.serving_output(outputs)
    
myswin = MySwin.from_pretrained(swinModel)
tf.saved_model.save(myswin, swin_EXPORT_PATH, signatures={
    "serving1": myswin.serving1,
    # "serving2": mygpt2.serving2
})

```

```

All model checkpoint layers were used when initializing MySwin.

All the layers of MySwin were initialized from the model checkpoint at microsoft/swin-tiny-patch4-window7-224.
If your task is similar to the task the model of the checkpoint was trained on, you can already use MySwin for predictions without further training.
---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
[<ipython-input-13-b219bb00369a>](https://localhost:8080/#) in <module>()
      1 myswin = MySwin.from_pretrained(swinModel)
      2 tf.saved_model.save(myswin, swin_EXPORT_PATH, signatures={
----> 3     "serving1": myswin.serving1,
      4     # "serving2": mygpt2.serving2
      5 })

14 frames
[/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py](https://localhost:8080/#) in autograph_handler(*args, **kwargs)
   1145           except Exception as e:  # pylint:disable=broad-except
   1146             if hasattr(e, "ag_error_metadata"):
-> 1147               raise e.ag_error_metadata.to_exception(e)
   1148             else:
   1149               raise

OperatorNotAllowedInGraphError: in user code:

    File "<ipython-input-11-84a42b1aca69>", line 10, in serving1  *
        outputs = self.call(pixel_values=inputs["pixel_values"])
    File "/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py", line 1426, in run_call_with_unpacked_inputs  *
        return func(self, **unpacked_inputs)
    File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 1439, in call  *
        outputs = self.swin(
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
        raise e.with_traceback(filtered_tb) from None

    OperatorNotAllowedInGraphError: Exception encountered when calling layer "swin" (type TFSwinMainLayer).
    
    in user code:
    
        File "/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py", line 1426, in run_call_with_unpacked_inputs  *
            return func(self, **unpacked_inputs)
        File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 1142, in call  *
            encoder_outputs = self.encoder(
        File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
    
        OperatorNotAllowedInGraphError: Exception encountered when calling layer "encoder" (type TFSwinEncoder).
        
        in user code:
        
            File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 906, in call  *
                layer_outputs = layer_module(
            File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
                raise e.with_traceback(filtered_tb) from None
        
            OperatorNotAllowedInGraphError: Exception encountered when calling layer "layers.0" (type TFSwinStage).
            
            in user code:
            
                File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 838, in call  *
                    layer_outputs = layer_module(
                File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
                    raise e.with_traceback(filtered_tb) from None
            
                OperatorNotAllowedInGraphError: Exception encountered when calling layer "blocks.0" (type TFSwinLayer).
                
                in user code:
                
                    File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 733, in call  *
                        self.set_shift_and_window_size(input_dimensions)
                    File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 672, in set_shift_and_window_size  *
                        if min(input_resolution) <= self.window_size:
                
                    OperatorNotAllowedInGraphError: using a `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:
                  • hidden_states=tf.Tensor(shape=(None, None, 96), dtype=float32)
                  • input_dimensions=('tf.Tensor(shape=(), dtype=int32)', 'tf.Tensor(shape=(), dtype=int32)')
                  • head_mask=None
                  • output_attentions=False
                  • training=False
            
            
            Call arguments received:
              • hidden_states=tf.Tensor(shape=(None, None, 96), dtype=float32)
              • input_dimensions=('tf.Tensor(shape=(), dtype=int32)', 'tf.Tensor(shape=(), dtype=int32)')
              • head_mask=None
              • output_attentions=False
              • training=False
        
        
        Call arguments received:
          • hidden_states=tf.Tensor(shape=(None, None, 96), dtype=float32)
          • input_dimensions=('tf.Tensor(shape=(), dtype=int32)', 'tf.Tensor(shape=(), dtype=int32)')
          • head_mask=['None', 'None', 'None', 'None']
          • output_attentions=False
          • output_hidden_states=False
          • return_dict=True
          • training=False
    
    
    Call arguments received:
      • self=tf.Tensor(shape=(None, None, None, None), dtype=float32)
      • pixel_values=None
      • bool_masked_pos=None
      • head_mask=None
      • output_attentions=False
      • output_hidden_states=False
      • return_dict=True
      • training=False


```


### Expected behavior

It is supposed to make a SavedModel but instead, I get this error mentioned above. The SavedModel is needed for TensorFlow Serving .

Comments:
cc @gante and @amyeroberts 

@amyeroberts is this related to what you've been investigating? If not, I can have a go at it :)

@gante Yep - I believe so. I've opened a PR here: https://github.com/huggingface/transformers/pull/18153

Hey @amyeroberts I updated my transformer code locally with the changes you made in the PR for swin but I still get the same result

```
---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
[<ipython-input-4-637c488e6341>](https://localhost:8080/#) in <module>()
----> 1 model.save_pretrained("test",saved_model=True)

2 frames
[/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py](https://localhost:8080/#) in autograph_handler(*args, **kwargs)
   1145           except Exception as e:  # pylint:disable=broad-except
   1146             if hasattr(e, "ag_error_metadata"):
-> 1147               raise e.ag_error_metadata.to_exception(e)
   1148             else:
   1149               raise

OperatorNotAllowedInGraphError: in user code:

    File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 979, in serving  *
        output = self.call(inputs)
    File "/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py", line 1457, in run_call_with_unpacked_inputs  *
        return func(self, **unpacked_inputs)
    File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 1470, in call  *
        outputs = self.swin(
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
        raise e.with_traceback(filtered_tb) from None

    OperatorNotAllowedInGraphError: Exception encountered when calling layer "swin" (type TFSwinMainLayer).
    
    in user code:
    
        File "/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py", line 1457, in run_call_with_unpacked_inputs  *
            return func(self, **unpacked_inputs)
        File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 1160, in call  *
            encoder_outputs = self.encoder(
        File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
    
        OperatorNotAllowedInGraphError: Exception encountered when calling layer "encoder" (type TFSwinEncoder).
        
        in user code:
        
            File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 905, in call  *
                layer_outputs = layer_module(
            File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
                raise e.with_traceback(filtered_tb) from None
        
            OperatorNotAllowedInGraphError: Exception encountered when calling layer "layers.0" (type TFSwinStage).
            
            in user code:
            
                File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 837, in call  *
                    layer_outputs = layer_module(
                File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
                    raise e.with_traceback(filtered_tb) from None
            
                OperatorNotAllowedInGraphError: Exception encountered when calling layer "blocks.0" (type TFSwinLayer).
                
                in user code:
                
                    File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 732, in call  *
                        self.set_shift_and_window_size(input_dimensions)
                    File "/usr/local/lib/python3.7/dist-packages/transformers/models/swin/modeling_tf_swin.py", line 671, in set_shift_and_window_size  *
                        if min(input_resolution) <= self.window_size:
                
                    OperatorNotAllowedInGraphError: using a `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:
                  • hidden_states=tf.Tensor(shape=(None, None, 96), dtype=float32)
                  • input_dimensions=('tf.Tensor(shape=(), dtype=int32)', 'tf.Tensor(shape=(), dtype=int32)')
                  • head_mask=None
                  • output_attentions=False
                  • training=False
            
            
            Call arguments received:
              • hidden_states=tf.Tensor(shape=(None, None, 96), dtype=float32)
              • input_dimensions=('tf.Tensor(shape=(), dtype=int32)', 'tf.Tensor(shape=(), dtype=int32)')
              • head_mask=None
              • output_attentions=False
              • training=False
        
        
        Call arguments received:
          • hidden_states=tf.Tensor(shape=(None, None, 96), dtype=float32)
          • input_dimensions=('tf.Tensor(shape=(), dtype=int32)', 'tf.Tensor(shape=(), dtype=int32)')
          • head_mask=['None', 'None', 'None', 'None']
          • output_attentions=False
          • output_hidden_states=False
          • return_dict=True
          • training=False
    
    
    Call arguments received:
      • self=tf.Tensor(shape=(None, None, None, None), dtype=float32)
      • pixel_values=None
      • bool_masked_pos=None
      • head_mask=None
      • output_attentions=False
      • output_hidden_states=False
      • return_dict=True
      • training=False
```

![image](https://user-images.githubusercontent.com/66001253/179504748-a1673b1d-0261-45c8-88e9-8319bf89d243.png)


OK @ahmedlone127. Thanks for letting me know. I'll dig into this some more.

Thanks 

Following merging of #18153 the reproduction snippet runs on main without error.