

TITLE
Cannot trace_module on models using model's generate function

COMMENTS
6

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

BODY
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.3.1
- Platform: Linux-5.8.14_1-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.6.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
@patrickvonplaten, @sshleifer
 -->

@patrickvonplaten, @sshleifer

## Information

Model I am using BART

The problem arises when using:
* [*] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [*] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. load any model that uses the generate function
2. try to trace it using trace_module

<!-- 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.-->

Can be easily reproduced with the following snippet:

```
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model = 'sshleifer/bart-tiny-random'
tokenizer = AutoTokenizer.from_pretrained(model)
sqgen_model = AutoModelForSeq2SeqLM.from_pretrained(model, torchscript=True)
sqgen_model.eval()
dummy_input = ' '.join('dummy' for dummy in range(512))
batch = tokenizer(
    [dummy_input], return_tensors='pt', truncation=True, padding='longest',
)
with torch.no_grad():
    traced_model = torch.jit.trace_module(  # type: ignore
        sqgen_model,
        {
            'forward': (batch.input_ids, batch.attention_mask),
            'generate': (batch.input_ids, batch.attention_mask),
        },
    )

```

It throws an error:

```
  File "/home/void/.miniconda3/envs/lexml/src/transformers/src/transformers/generation_utils.py", line 288, in generate
    assert isinstance(max_length, int) and max_length > 0, "`max_length` should be a strictly positive integer."
AssertionError: `max_length` should be a strictly positive integer.
```

obviously because the generate function's second argument is supposed to be max_length and not attention_mask

## Expected behavior
Should be able to trace models that use the generate function.

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