From 4eff7009ca1241248f703de5e92d4c5de7d4563b Mon Sep 17 00:00:00 2001
From: shibing624 <shibing624@126.com>
Date: Wed, 20 Sep 2023 11:00:11 +0800
Subject: [PATCH] move model to device.

---
 text2vec/sentence_model.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/text2vec/sentence_model.py b/text2vec/sentence_model.py
index 3273f20..605aee5 100644
--- 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: