add comment

sbintuitions · Aug 8, 2024 · 39f98a3 · 39f98a3
1 parent 4261cf5
commit 39f98a3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/jmteb/embedders/data_parallel_sbert_embedder.py b/src/jmteb/embedders/data_parallel_sbert_embedder.py
@@ -17,6 +17,7 @@
 
 
 class DPSentenceTransformer(SentenceTransformer):
+    """SentenceBERT with pytorch torch.nn.DataParallel"""
 
     def __init__(self, sbert_model: SentenceTransformer):
         super(DPSentenceTransformer, self).__init__()
@@ -209,6 +210,7 @@ def _encode_with_auto_batch_size(batch_size, self, text, prefix):
                     batch_size=batch_size,
                     normalize_embeddings=self.normalize_embeddings,
                 )
+
                 self.batch_size = batch_size
                 return out