Cannot load existing model

[zhao-zilong]

Hi, I tried to load a trained model as follows:

model = GReaT(llm='distilgpt2', batch_size=32, epochs=1)
model.load_from_dir("trainer_great/checkpoint-76000")

I checked that I have the files there, but it gives me following errors:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[27], line 1
----> 1 model.load_from_dir("trainer_great/checkpoint-76000")

Cell In[2], line 254, in GReaT.load_from_dir(cls, path)
    251     attributes = json.load(f)
    253 # Create new be_great model instance
--> 254 great = cls(attributes["llm"])
    256 # Set all attributes
    257 for k, v in attributes.items():

KeyError: 'llm'

Thanks for your attention

[zhao-zilong]

oh, I actually found the problem, if we directly use the library to train the model, the automatically generated checkpoint does not contain the model.pt file, we need to explicitly call the save() function, then the model is saved.