The generation input shape and the output shape from the official scripts are completely different  for the TFLite model 
Author: generic-matrix
Pull request: no

### System Info

- `transformers` version: 2.3.0
- Platform: Linux-5.10.133+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.15
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): 2.9.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

and 

- `transformers` version: 4.25.1
- Platform: Linux-5.10.133+-x86_64-with-glibc2.27
- Python version: 3.8.16
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.13.0+cu116 (False)
- Tensorflow version (GPU?): 2.9.2 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@patric @anton-l  @sanchit-gandhi  @Rocketknight1 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Upon cloning the gpt2.py from the [link ](https://github.com/huggingface/tflite-android-transformers/tree/master/models_generation) renders input shape as [3 5] instead of [1 64]

```
import transformers
import tensorflow
print(transformers.__version__)
print(tensorflow.__version__)
```

```
2.3.0
2.2.0
```





```
!wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-64.tflite

import numpy as np
import tensorflow as tf

tflite_model_path = 'gpt2-64.tflite'
# Load the TFLite model and allocate tensors
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()

# Get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']

#print the output
input_data = np.array(np.random.random_sample((input_shape)), dtype=np.int32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data.shape)
print(input_shape)

```

```
>(1, 64, 50257)
>[ 1 64]
```





```
import tensorflow as tf
from transformers import TFGPT2LMHeadModel
import numpy as np

model = TFGPT2LMHeadModel.from_pretrained('gpt2') # or 'distilgpt2'

input_spec = tf.TensorSpec([1, 64], tf.int32)
model._set_inputs(input_spec, training=False)

converter = tf.lite.TFLiteConverter.from_keras_model(model)

# For FP16 quantization:
# converter.optimizations = [tf.lite.Optimize.DEFAULT]
# converter.target_spec.supported_types = [tf.float16]

tflite_model = converter.convert()

open("gpt2-64-2.tflite", "wb").write(tflite_model)

tflite_model_path = 'gpt2-64-2.tflite'
# Load the TFLite model and allocate tensors
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()

# Get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']

#print the output
input_data = np.array(np.random.random_sample((input_shape)), dtype=np.int32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data.shape)
print(input_shape)

```

```
>(3, 5, 50257)
>[3 5]

```
**Expected 
>(1, 64, 50257)
>[ 1 64]**

```
import transformers
import tensorflow
print(transformers.__version__)
print(tensorflow.__version__)

>4.25.1
> 2.9.2
```

```
!wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-64.tflite

import numpy as np
import tensorflow as tf

tflite_model_path = 'gpt2-64.tflite'
# Load the TFLite model and allocate tensors
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()

# Get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']

#print the output
input_data = np.array(np.random.random_sample((input_shape)), dtype=np.int32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data.shape)
print(input_shape)

```

```
>(1, 64, 50257)
> [ 1 64]
```


```
import tensorflow as tf
from transformers import TFGPT2LMHeadModel
import numpy as np

model = TFGPT2LMHeadModel.from_pretrained('gpt2') # or 'distilgpt2'

input_spec = tf.TensorSpec([1, 64], tf.int32)
model._set_inputs(input_spec, training=False)

converter = tf.lite.TFLiteConverter.from_keras_model(model)

# For FP16 quantization:
# converter.optimizations = [tf.lite.Optimize.DEFAULT]
# converter.target_spec.supported_types = [tf.float16]

tflite_model = converter.convert()

open("gpt2-64-2.tflite", "wb").write(tflite_model)

tflite_model_path = 'gpt2-64-2.tflite'
# Load the TFLite model and allocate tensors
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()

# Get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

input_shape = input_details[0]['shape']

#print the output
input_data = np.array(np.random.random_sample((input_shape)), dtype=np.int32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data.shape)
print(input_shape)
```

```
>(2, 1, 12, 1, 64)
>[1 1]
```
**Expected 
>(1, 64, 50257)
>[ 1 64]**


How can we fix the same ?

### Expected behavior

Expected inputshape is [ 1 64] and output shape is (1, 64, 50257) 

Comments:
cc @gante and @Rocketknight1 

@sgugger @gante @Rocketknight1 any update on the same ?

I even tried the same with the collab below taking latest version of transformers tensorflow into consideration , I get the same issue as above. When we import the new model with different output shape onto the android project (gpt2) , I get the issue as below :

```
E/AndroidRuntime: FATAL EXCEPTION: main
    Process: co.huggingface.android_transformers.gpt2, PID: 17293
    java.lang.IllegalArgumentException: ByteBuffer is not a valid flatbuffer model
        at org.tensorflow.lite.NativeInterpreterWrapper.createModelWithBuffer(Native Method)
        at org.tensorflow.lite.NativeInterpreterWrapper.<init>(NativeInterpreterWrapper.java:60)
        at org.tensorflow.lite.Interpreter.<init>(Interpreter.java:224)
        at co.huggingface.android_transformers.gpt2.ml.GPT2Client$loadModel$2.invokeSuspend(GPT2Client.kt:138)
        at kotlin.coroutines.jvm.internal.BaseContinuationImpl.resumeWith(ContinuationImpl.kt:33)
        at kotlinx.coroutines.DispatchedTask.run(Dispatched.kt:241)
        at kotlinx.coroutines.scheduling.CoroutineScheduler.runSafely(CoroutineScheduler.kt:594)
        at kotlinx.coroutines.scheduling.CoroutineScheduler.access$runSafely(CoroutineScheduler.kt:60)
        at kotlinx.coroutines.scheduling.CoroutineScheduler$Worker.run(CoroutineScheduler.kt:740)
```

https://gist.github.com/sudo-carson/158d9b9e7208e42977b08d966f3f4989



Hello @sgugger @gante @Rocketknight1 

The issue has been fixed, the below code can be used.

The output should be keras_output.logits as in the code below 

```
model = TFGPT2LMHeadModel.from_pretrained('gpt2') # or 'distilgpt2'
input = tf.keras.Input([ 64 ], batch_size=1, dtype=tf.int32)
keras_output = model(input, training=False)
model = tf.keras.Model(input, keras_output.logits)
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# For FP16 quantization:
# converter.optimizations = [tf.lite.Optimize.DEFAULT]
# converter.target_spec.supported_types = [tf.float16]

tflite_model = converter.convert()

open("model.tflite", "wb").write(tflite_model)
```

Hey @generic-matrix 👋 

Thank you for raising the issue and for reporting its fix as well <3 For context, we haven't been checking whether our models are supported by TFLite. That's something we plan to rectify over this year, with notebooks and demos as well!

(closing as it is fixed)