From b89517fde46068745643c015b48817f312f49e24 Mon Sep 17 00:00:00 2001
From: Antoine Chaffin <38869395+NohTow@users.noreply.github.com>
Date: Thu, 17 Oct 2024 17:01:49 +0200
Subject: [PATCH] Fixing issue with ST 3.1 and bumping version (#65)

---
 pylate/models/colbert.py | 6 +++---
 setup.py                 | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pylate/models/colbert.py b/pylate/models/colbert.py
index a3e4418..ca64cbb 100644
--- a/pylate/models/colbert.py
+++ b/pylate/models/colbert.py
@@ -301,7 +301,7 @@ def __init__(
             )
             self[1] = Dense.from_sentence_transformers(dense=self[1])
         else:
-            logger.warning("Pylate model loaded successfully.")
+            logger.warning("PyLate model loaded successfully.")
 
         # Ensure all tensors in the model are of the same dtype as the first tensor
         try:
@@ -1150,7 +1150,7 @@ def _load_sbert_model(
         config_kwargs: dict | None = None,
     ) -> list[nn.Module]:
         """Create a Sentence Transformer model from a model name or path."""
-        modules = super()._load_sbert_model(
+        modules, module_kwargs = super()._load_sbert_model(
             model_name_or_path=model_name_or_path,
             token=token,
             cache_folder=cache_folder,
@@ -1196,4 +1196,4 @@ def _load_sbert_model(
             for module in modules.values()
             if isinstance(module, Transformer)
             or isinstance(module, DenseSentenceTransformer)
-        ]
+        ], module_kwargs
diff --git a/setup.py b/setup.py
index 278a4f4..75b2e74 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,7 @@
     long_description = fh.read()
 
 base_packages = [
-    "sentence-transformers == 3.0.1",
+    "sentence-transformers == 3.2.0",
     "datasets >= 2.20.0",
     "accelerate >= 0.31.0",
     "voyager >= 2.0.9",