

TITLE
GeneratorExp aren't supported by torch.jit.script when I try to export a previously trained model  'google/vit-base-patch16-224-in21k'.

COMMENTS
3

REACTIONS
+1: 1
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## Environment info
- `transformers` version: 4.15.0
- Platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.12
- PyTorch version (GPU?): 1.10.0+cu111 (True)
- Tensorflow version (GPU?): 2.7.0 (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

Models:

ViTModel

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Vision: @NielsRogge, @sgugger

Documentation: @sgugger

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.



## Information

GeneratorExp aren't supported by torch.jit.script when I try to export a previously trained model  'google/vit-base-patch16-224-in21k'.

Model I am using (ViTModel):

The problem arises when using:
* [X] my own modified scripts: (give details below)
* 
model_x = ViTForImageClassification.from_pretrained(
    'google/vit-base-patch16-224-in21k',
    num_labels=len(label2id),
    label2id=label2id,
    id2label=id2label
)
model_scripted = torch.jit.script(model_x) # Export to TorchScript

---------------------------------------------------------------------------
UnsupportedNodeError                      Traceback (most recent call last)
<ipython-input-12-bc467d8ea1c0> in <module>()
      6     id2label=id2label
      7 )
----> 8 model_scripted = torch.jit.script(model_x) # Export to TorchScript
      9 model_scripted.save('model_scripted.pt') # Save

14 frames
/usr/local/lib/python3.7/dist-packages/torch/jit/frontend.py in __call__(self, ctx, node)
    284         method = getattr(self, 'build_' + node.__class__.__name__, None)
    285         if method is None:
--> 286             raise UnsupportedNodeError(ctx, node)
    287         return method(ctx, node)
    288 

UnsupportedNodeError: GeneratorExp aren't supported:
  File "/usr/local/lib/python3.7/dist-packages/transformers/modeling_utils.py", line 987
        activations".
        """
        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
                   ~ <--- HERE

## To reproduce

Steps to reproduce the behavior:

1.  from transformers import ViTForImageClassification
2.   Instantiate a previously created mode  'google/vit-base-patch16-224-in21k'  using ViTForImageClassification.from_pretrained() API.
3. Try invoking  torch.jit.script(model_x) and you will see the error.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
