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

@@ -20,11 +20,16 @@ from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM
 from transformers.tokenization_utils_base import BatchEncoding
 from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 
+import logging
+
+logger = logging.getLogger(__name__)
+print("Running code of HF Model")
+
 # pyright: reportUnknownMemberType=false
 # 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")
 )
@@ -126,9 +131,11 @@ def predict(
         query_documents = [[sentence[0], sentence[1]] for sentence in sentences]
 
     if not hasattr(self, "inner_model"):
+        logger.info(f"Memory reserved [Within Model File] Before Loading Model: {torch.cuda.memory_reserved()}")
         self.inner_tokenizer, self.inner_model = load_model(global_device)
-        self.inner_model.gradient_checkpointing_enable()
+        logger.info(f"Memory reserved [Within Model File] After Loading Model: {torch.cuda.memory_reserved()}")
         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 +179,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()