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")
 )
@@ -127,8 +127,8 @@ def predict(
 
     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_model.gradient_checkpointing_disable()
         self.inner_yes_token_id = self.inner_tokenizer.encode(
             "Yes", add_special_tokens=False
         )[0]
@@ -172,7 +172,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()