Update helpers/required_classes.py

helpers/required_classes.py

@@ -24,14 +24,15 @@ class BertEmbedder:
         self.embedder.to(self.device)
 
     def __call__(self, text: str):
-        encoded_input = self.tokenizer(text,
-                                       return_tensors='pt',
-                                       max_length=self.max_length,
-                                       padding=True,
-                                       truncation=True).to(self.device)
-        model_output = self.embedder(**encoded_input)
-        text_embed = model_output.pooler_output[0].cpu()
-        return text_embed
+        with torch.no_grad():
+            encoded_input = self.tokenizer(text,
+                                           return_tensors='pt',
+                                           max_length=self.max_length,
+                                           padding=True,
+                                           truncation=True).to(self.device)
+            model_output = self.embedder(**encoded_input)
+            text_embed = model_output.pooler_output[0].cpu()
+            return text_embed
 
     def batch_predict(self, texts: List[str]):
         encoded_input = self.tokenizer(texts,