added rnn.pack_padded_sequence argument cpu conversion (#1420)

UKPLab · Dec 18, 2023 · c006921 · c006921
1 parent 6d11984
commit c006921
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sentence_transformers/models/LSTM.py b/sentence_transformers/models/LSTM.py
@@ -29,7 +29,7 @@ def forward(self, features):
         token_embeddings = features['token_embeddings']
         sentence_lengths = torch.clamp(features['sentence_lengths'], min=1)
 
-        packed = nn.utils.rnn.pack_padded_sequence(token_embeddings, sentence_lengths, batch_first=True, enforce_sorted=False)
+        packed = nn.utils.rnn.pack_padded_sequence(token_embeddings, sentence_lengths.cpu(), batch_first=True, enforce_sorted=False)
         packed = self.encoder(packed)
         unpack = nn.utils.rnn.pad_packed_sequence(packed[0], batch_first=True)[0]
         features.update({'token_embeddings': unpack})