Saving trained model for testing

I want to have two separate programs, one that trains my model, saves the trained parameters, and the other to use those trained parameters to make predictions. I’ve tried placing:

torch.save(model, 'model_best.pt')

in one program. And in the testing program:

model = torch.load('model_best.pt')

Then I use the model to predict, e.g.:

output = model(x).

However, every time I call the torch.load, it seems to force me to retrain my model - it somehow calls my training program. Can anyone explain this behavior, and how you are supposed to properly do this?

Thanks

Try to save and load the state_dict. Serialization semantics

I tried that method too - it gives me the same behavior. It will also throw this afterwards:

    outputs = model(x)

TypeError: 'NoneType' object is not callable

Seems like whenever I try to import my file with the nn.Module class, just due to the import line, it will try to retrain my model…