Save and re-load a trained prompt model

[agrimaseth]

Currently, I am saving the model using the following: torch.save(prompt_model.state_dict(), PATH)

How can we load this back to test performance on other data? Pytorch tutorial says to use the following. model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH)) model.eval()

What would TheModelClass be? An example would be really appreciated.

[agrimaseth]

@ShengdingHu could you please help with this? Using the above load command doesn't work.