avoid gradients tracking, for faster inference and less memory consumption

aditya0by0 · aditya0by0 · commit 3d7d91855bd0 · 2025-11-22T21:04:55.000+01:00
diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py
@@ -22,6 +22,8 @@ def __init__(
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model = self.init_model(ckpt_path=ckpt_path)
+        self.model.eval()
+
         self.target_labels = [
             line.strip() for line in open(target_labels_path, encoding="utf-8")
         ]
@@ -32,6 +34,7 @@ def init_model(self, ckpt_path: str, **kwargs):
             "Model initialization must be implemented in subclasses."
         )
 
+    @torch.inference_mode()
     def calculate_results(self, batch):
         collator = self.reader_cls.COLLATOR()
         dat = self.model._process_batch(collator(batch).to(self.device), 0)

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,8 @@ def __init__(`
`22`	`22`
`23`	`23`	`self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")`
`24`	`24`	`self.model = self.init_model(ckpt_path=ckpt_path)`
	`25`	`+ self.model.eval()`
	`26`	`+`
`25`	`27`	`self.target_labels = [`
`26`	`28`	`line.strip() for line in open(target_labels_path, encoding="utf-8")`
`27`	`29`	`]`
`@@ -32,6 +34,7 @@ def init_model(self, ckpt_path: str, **kwargs):`
`32`	`34`	`"Model initialization must be implemented in subclasses."`
`33`	`35`	`)`
`34`	`36`
	`37`	`+ @torch.inference_mode()`
`35`	`38`	`def calculate_results(self, batch):`
`36`	`39`	`collator = self.reader_cls.COLLATOR()`
`37`	`40`	`dat = self.model._process_batch(collator(batch).to(self.device), 0)`