Aitextgen doesn't generate text when generate with prompt

Question

Im using aitextgen on colab, done finetuning. But when i generate with prompt, like this:

ai.generate(n=1,
            batch_size=50,
            prompt="Hello",
            max_length=100,
            top_k=40,
            temperature=2.0,
            top_p=0.9,
            no_repeat_ngram_size=4
           )

Then most of the time it doesn't generate anything at all, just send out the prompt (in this case is 'Hello') as output and nothing else, it only starting to generate when i increase the tempature to >2.0. The text generation without prompt work normally, regarless of the tempature:

ai.generate()

How do i fix this? Is this the code bugs, or is it something related to my model?

Here the code i used on colab:

#Install requirements and import
!pip install -qq pytorch-lightning==1.8.4 transformers==4.21.3
!pip install git+https://github.com/llimllib/aitextgen@fix_tpu_available

import logging
logging.basicConfig(
        format="%(asctime)s — %(levelname)s — %(name)s — %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO
    )

from aitextgen import aitextgen
from aitextgen.colab import mount_gdrive, copy_file_from_gdrive

#Load gpt-neo model
ai = aitextgen(model="EleutherAI/gpt-neo-125m", to_gpu=True)

#Mount drive
mount_gdrive()

#Upload dataset file and run block
file_name_1 = "dataset.txt"

#Start finetune
ai.train(file_name_1,
         line_by_line=False,
         from_cache=False,
         num_steps=500,
         generate_every=50,
         save_every=50,
         save_gdrive=False,
         learning_rate=1e-4,
         fp16=False,
         batch_size=1,
         )

#Load trained model
ai = aitextgen(model_folder="trained_model", to_gpu=True)

#No prompt generate
ai.generate()

#Generate with prompt
ai.generate(n=1,
            batch_size=50,
            prompt="Hello",
            max_length=100,
            top_k=40,
            temperature=2.0,
            top_p=0.9,
            no_repeat_ngram_size=4
           )

I tried to further increase the tempature to see if it generate more, but it still doesnt generate some of the time