Storing and restoring GPT-J model

Hi,
I am trying to save and later restore the GPT-J models as shown below. But I am getting an error

AttributeError: ‘BaseModelOutputWithPast’ object has no attribute ‘logits’

Any help with this will be much appreciated

Saving the model

from transformers import GPTJModel, GPTJConfig

print('---start---')
config = GPTJConfig.from_pretrained("EleutherAI/gpt-j-6B")
model = GPTJModel(config)
model.save_pretrained("my-gpt-j-model")
from transformers import AutoTokenizer

print('---start---')
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
tokenizer.save_pretrained("my-gpt-j-tokenizer")

Restoring and using the model

This piece of code is throwing the error

from transformers import AutoTokenizer
from transformers import GPTJModel

#gpt_config = GPTJConfig.from_pretrained("my-gpt-j-model")
gpt_model = GPTJModel.from_pretrained("my-gpt-j-model")
t=gpt_model.eval()
print(t)
input_text = "Hello my name is Blake and"

tokenizer = AutoTokenizer.from_pretrained("my-gpt-j-tokenizer")
input_ids = tokenizer.encode(str(input_text),return_tensors='pt')

output=gpt_model.generate(
    input_ids,
    do_sample=True,
    max_length=20,
    top_p=0.7,
    top_k=0,
    temperature=1.0,
)
print(tokenizer.decode(output[0],skip_special_tokens=True))