Update model.py

model.py

@@ -150,7 +150,7 @@ class Trainer:
     
     def predict(self, inputs):
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.load_state_dict(torch.load('best_model.pth', weights_only=True))
+        self.model.load_state_dict(torch.load('best_model.pth', weights_only=True, map_location=torch.device('cpu'))
         self.model.to(device)
         self.model.eval()
         inputs = inputs.to(device)
@@ -158,11 +158,11 @@ class Trainer:
             outputs = self.model(inputs)
         return outputs
 
-if __name__ == "__main__":
-    trainer = Trainer()
+# if __name__ == "__main__":
+#     trainer = Trainer()
     
-    # Test the model
-    image_path = '6082308423334085331.jpg'
-    image_tensor = trainer.image_to_tensor(image_path)
-    prediction = trainer.predict(image_tensor)
-    print(f"Predicted Beauty Score: {prediction.item() * 100}")
+#     # Test the model
+#     image_path = '6082308423334085331.jpg'
+#     image_tensor = trainer.image_to_tensor(image_path)
+#     prediction = trainer.predict(image_tensor)
+#     print(f"Predicted Beauty Score: {prediction.item() * 100}")