Spaces:

waseoke
/

shortpingoo

Sleeping

waseoke commited on Dec 7, 2024

Commit

2642269

verified ·

1 Parent(s): 067bb59

Update calculate_cosine_similarity.py

Files changed (1) hide show

calculate_cosine_similarity.py CHANGED Viewed

@@ -48,7 +48,7 @@ def load_trained_model(model_path="product_model.pth"):
         torch.nn.ReLU(),
         torch.nn.Linear(256, 128),
     )
-    model.load_state_dict(torch.load(model_path))
     model.eval()  # 평가 모드
     return model

         torch.nn.ReLU(),
         torch.nn.Linear(256, 128),
     )
+    model.load_state_dict(torch.load(model_path, weights_only=True))
     model.eval()  # 평가 모드
     return model