How to perform inference by loading the saved best_model?

[gy2256]

Hello,

I finished training my model and saved the best_model using

trainer.train()
torch.save(trainer.best_model, 'best_model.pth')

Now, in my other file, I want to perform inference given the input with the trained model. What is the proper way to do it?

Do I need to create a new Problem() and then load the pth file?

self.model = problem([dynamics_model], loss)
self.model.load_state_dict(torch.load('best_model.pth'), strict=False)
self.model.eval()

output = self.model(input)

However, this requires me to define create theproblem([dynamics_model], loss)by defining [dynamics_model] and loss again, but I don't need to calculate the loss during deployment. Is there a way for to save the trained MLP directly as .pth file?

[gy2256]

Solved by using torch.save(fx, "model_name.pth")