From b34bba0b9efa1e7d605463bd8354a2030d11c6f1 Mon Sep 17 00:00:00 2001
From: Antoine Chaffin <38869395+NohTow@users.noreply.github.com>
Date: Tue, 22 Oct 2024 13:28:25 +0200
Subject: [PATCH] Correctly setting safetensor path if the stanford model is a
 local folder (#68)

---
 pylate/models/Dense.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pylate/models/Dense.py b/pylate/models/Dense.py
index fb37256..b9f3575 100644
--- a/pylate/models/Dense.py
+++ b/pylate/models/Dense.py
@@ -121,6 +121,9 @@ def from_stanford_weights(
                 token=token,
                 use_auth_token=use_auth_token,
             )
+        # If the model a local folder, load the safetensor
+        else:
+            model_name_or_path = os.path.join(model_name_or_path, "model.safetensors")
         with safe_open(model_name_or_path, framework="pt", device="cpu") as f:
             state_dict = {"linear.weight": f.get_tensor("linear.weight")}