fix: no grad

MinishLab · Dec 24, 2024 · 4713bfa · 4713bfa
1 parent c4ba272
commit 4713bfa
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py
@@ -48,6 +48,7 @@ def predict(self, texts: list[str]) -> list[str]:
 
         return [self.classes[idx] for idx in logits.argmax(1)]
 
+    @torch.no_grad()
     def _predict(self, texts: list[str]) -> torch.Tensor:
         input_ids = self.tokenize(texts)
         vectors, _ = self.forward(input_ids)