# Issues with model.generate()

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fissues-with-model-generate%2F39023%3E%5C
**Category:** Models\
**Created:** [May 8, 2023, 1:18am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fissues-with-model-generate%2F39023 "2023-05-08T01:18:53Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Shilz](https://avatars.discourse-cdn.com/v4/letter/s/34f0e0/32.png) [@Shilz](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FShilz)%5C
**Post date:** [May 8, 2023, 1:18am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fissues-with-model-generate%2F39023%2F1 "2023-05-08T01:18:53Z")

</div>

Hi, I am trying to train distilgpt-2 with a custom head . After training for a while when I try to generate text using model.generate(), I get an error “InvalidArgumentError: {{function\_node \__wrapped\_\_Reshape\_device_/job:localhost/replica:0/task:0/device:GPU:0}} Input to reshape is a tensor with 30720 values, but the requested shape requires a multiple of 76800 [Op:Reshape]”  
My LM head class is as follows 🙂

```auto
class policy_rl(TFGPT2LMHeadModel):
  def __init__ (self,config):
    super(). __init__ (config)
    self.transformer = TFGPT2Model(config)
    self.lm_head = tf.keras.layers.Dense(config.vocab_size, input_shape=(config.n_embd,), use_bias=False)
    

  def call(self, input_ids=None, past_key_values=None, attention_mask=None, token_type_ids=None, position_ids=None,
             head_mask=None, inputs_embeds=None,output_attentions=None,output_hidden_states=None,
           return_dict= None, labels=None, use_cache=True, get_hidden=False):
    transformer_outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask,
                                               token_type_ids=token_type_ids, position_ids=position_ids,
                                               head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache)
    hidden_states = transformer_outputs[0]

    if get_hidden:
      return hidden_states
    logits = self.lm_head(hidden_states)
    loss = None
    #labels = input_ids
    if labels is not None:
      # shift labels to the left and cut last logit token
      print('inside labels')
      logits = logits[:, :-1]
      labels = labels[:, 1:]
      loss = self.compute_loss(labels, logits)  
    #if not inputs["return_dict"]:
    output = (logits,) + transformer_outputs[1:]
    # return ((loss,) + output) if loss is not None else output 
    return TFCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,)
  def serving_output(self, output):
    pkv = tf.convert_to_tensor(output.past_key_values) if self.config.use_cache else None
    hs = tf.convert_to_tensor(output.hidden_states) if self.config.output_hidden_states else None
    attns = tf.convert_to_tensor(output.attentions) if self.config.output_attentions else None
    loss = tf.convert_to_tensor(output.loss)

    return TFCausalLMOutputWithPast(loss=loss,logits=output.logits, past_key_values=pkv, hidden_states=hs, attentions=attns)

```

I am trying to execute this 🙂

```auto
text = 'The fox jumped'
input_ids = tokenizer.encode(text, return_tensors='tf')
beam_output = model.generate(
  input_ids,
  past_key_values = None,
  max_length = 50,
  num_beams = 5,
  temperature = 0.7,
  no_repeat_ngram_size=2,
  num_return_sequences=2
)

```

I tried using the below as well :

```auto
output_model = model.generate(input_ids[0],max_length=10) # do greedy decoding
print(f"Generated: {tokenizer.decode(output_model, skip_special_tokens=True)}")

```

This gives the error :  
InvalidArgumentError: {{function\_node \__wrapped\_\_StridedSlice\_device_/job:localhost/replica:0/task:0/device:GPU:0}} Index out of range using input dim 1; input has only 1 dims [Op:StridedSlice] name: strided\_slice/

Please can anyone help me with this issue.

---

<div class="post-metadata">

**Author:** ![Shilz](https://avatars.discourse-cdn.com/v4/letter/s/34f0e0/32.png) [@Shilz](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FShilz)%5C
**Post date:** [May 8, 2023, 11:04pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fissues-with-model-generate%2F39023%2F2 "2023-05-08T23:04:20Z")

</div>

@Rocketknight1 I saw your comment on [Fine-Tuning TFGPT2LMHeadModel / What to pass to fit · Issue #11507 · huggingface/transformers · GitHub](https://github.com/huggingface/transformers/issues/11507). if you could please have a look.