Hi,
I am saving the weights of torchvision.models.resnet50 model after training it on my custom dataset using below code:
best_model_weights = copy.deepcopy(self.model.state_dict())
torch.save(best_model_weights, os.path.join(self.save_model_attributes_path, f"best_weights.pth"))
I am saving self.model.state_dict() correctly instead of the function self.model.state_dict
When I try to load the custom weights using the following code:
self.current_model = torchvision.models.resnet50
current_model = self.current_model.load_state_dict(torch.load(self.best_weights_path, device))
I am getting the following error:
Traceback (most recent call last):
File “C:\Macs Lab\imageClassification\prediction_model.py”, line 113, in
main()
File “C:\Macs Lab\imageClassification\prediction_model.py”, line 109, in main
best_weights_path=best_weights_path).classify_images_and_save_predictions()
File “C:\Macs Lab\imageClassification\prediction_model.py”, line 94, in classify_images_and_save_predictions
predictions_dict = self.predict_images()
File “C:\Macs Lab\imageClassification\prediction_model.py”, line 74, in predict_images
current_model = self.current_model.load_state_dict(copy.deepcopy(torch.load(self.best_weights_path, device)))
AttributeError: ‘function’ object has no attribute ‘load_state_dict’