move model to device.

shibing624 · Sep 20, 2023 · 4eff700 · 4eff700
1 parent abb6a07
commit 4eff700
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/text2vec/sentence_model.py b/text2vec/sentence_model.py
@@ -178,6 +178,8 @@ def encode(
         self.bert.eval()
         if device is None:
             device = self.device
+        self.bert.to(device)
+
         if max_seq_length is None:
             max_seq_length = self.max_seq_length
         if convert_to_tensor: