Use torch.inference_mode() and disable gradient checkpointing

config.json

@@ -64,5 +64,8 @@
   "transformers_version": "4.57.1",
   "use_cache": true,
   "use_sliding_window": false,
-  "vocab_size": 151936
+  "vocab_size": 151936,
+  "auto_map": {
+      "AutoConfig": "modeling_zeranker.ZEConfig"
+  }
 }

modeling_zeranker.py

@@ -24,7 +24,7 @@ from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 # pyright: reportUnknownVariableType=false
 
 MODEL_PATH = "zeroentropy/zerank-2"
-PER_DEVICE_BATCH_SIZE_TOKENS = 15_000
+PER_DEVICE_BATCH_SIZE_TOKENS = 10_000
 global_device = (
     torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 )
@@ -125,13 +125,7 @@ def predict(
             raise ValueError("query_documents or sentences must be provided")
         query_documents = [[sentence[0], sentence[1]] for sentence in sentences]
 
-    if not hasattr(self, "inner_model"):
-        self.inner_tokenizer, self.inner_model = load_model(global_device)
-        self.inner_model.gradient_checkpointing_enable()
-        self.inner_model.eval()
-        self.inner_yes_token_id = self.inner_tokenizer.encode(
-            "Yes", add_special_tokens=False
-        )[0]
+
 
     model = self.inner_model
     tokenizer = self.inner_tokenizer
@@ -172,7 +166,8 @@ def predict(
         batch_inputs = batch_inputs.to(global_device)
 
         try:
-            outputs = model(**batch_inputs, use_cache=False)
+            with torch.inference_mode():
+                outputs = model(**batch_inputs, use_cache=False)
         except torch.OutOfMemoryError:
             print(f"GPU OOM! {torch.cuda.memory_reserved()}")
             torch.cuda.empty_cache()
@@ -207,6 +202,22 @@ def to_device(self: _CE, new_device: torch.device) -> None:
     global_device = new_device
 
 
+
+original_init = _CE.__init__
+
+
+def _new_init(self: _CE, *args: Any, **kwargs: Any) -> None:
+    original_init(self, *args, **kwargs)
+    self.inner_tokenizer, self.inner_model = load_model(global_device)
+    self.inner_model.eval()
+    self.inner_model.gradient_checkpointing_disable()
+    self.inner_yes_token_id = self.inner_tokenizer.encode(
+        "Yes", add_special_tokens=False
+    )[0]
+
+
+_CE.__init__ = _new_init
+
 _CE.predict = predict
 
 from transformers import Qwen3Config