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

@@ -1,9 +1,8 @@
 from sentence_transformers import CrossEncoder as _CE
 
 import math
+import logging
 from typing import cast, Any
-import types
-
 
 import torch
 from transformers.configuration_utils import PretrainedConfig
@@ -23,8 +22,10 @@ from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 # pyright: reportUnknownMemberType=false
 # pyright: reportUnknownVariableType=false
 
+logger = logging.getLogger(__name__)
+
 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")
 )
@@ -74,9 +75,12 @@ def load_model(
     if device is None:
         device = global_device
 
+    logger.info(f"Loading model from {MODEL_PATH} on device: {device}")
+
     config = AutoConfig.from_pretrained(MODEL_PATH)
     assert isinstance(config, PretrainedConfig)
 
+    logger.info(f"Loading model with config type: {config.model_type}")
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_PATH,
         torch_dtype="auto",
@@ -93,6 +97,7 @@ def load_model(
         | Qwen3ForCausalLM,
     )
 
+    logger.info("Loading tokenizer")
     tokenizer = cast(
         AutoTokenizer,
         AutoTokenizer.from_pretrained(
@@ -105,6 +110,7 @@ def load_model(
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
+    logger.info("Model and tokenizer loaded successfully")
     return tokenizer, model
 
 
@@ -125,13 +131,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]
+    logger.info(f"Starting prediction for {len(query_documents)} query-document pairs")
 
     model = self.inner_model
     tokenizer = self.inner_tokenizer
@@ -161,9 +161,12 @@ def predict(
         batches[-1].append((query, document))
         max_length = max(max_length, 20 + len(query) + len(document))
 
+    logger.info(f"Created {len(batches)} batches for inference")
+
     # Inference all of the document batches
     all_logits: list[float] = []
-    for batch in batches:
+    for batch_idx, batch in enumerate(batches):
+        logger.debug(f"Processing batch {batch_idx + 1}/{len(batches)} with {len(batch)} pairs")
         batch_inputs = format_pointwise_datapoints(
             tokenizer,
             batch,
@@ -172,11 +175,12 @@ 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()}")
+            logger.warning(f"GPU OOM! Memory reserved: {torch.cuda.memory_reserved()}")
             torch.cuda.empty_cache()
-            print(f"GPU After OOM Cache Clear: {torch.cuda.memory_reserved()}")
+            logger.info(f"GPU cache cleared. Memory reserved: {torch.cuda.memory_reserved()}")
             outputs = model(**batch_inputs, use_cache=False)
 
         # Extract the logits
@@ -199,18 +203,31 @@ def predict(
     # Unsort by indices
     scores = [score for _, score in sorted(zip(permutation, scores, strict=True))]
 
+    logger.info(f"Prediction complete. Generated {len(scores)} scores")
     return scores
 
 
 def to_device(self: _CE, new_device: torch.device) -> None:
     global global_device
+    logger.info(f"Changing device from {global_device} to {new_device}")
     global_device = new_device
 
+    # Load the model now since __init__ patching doesn't work due to timing
+    # (CrossEncoder instance is created before this module is loaded)
+    if not hasattr(self, "inner_model"):
+        logger.info("Loading model during device setup (eager loading)")
+        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]
+        logger.info(f"Model loaded successfully. Yes token ID: {self.inner_yes_token_id}")
+
 
 _CE.predict = predict
+_CE.to = to_device
 
 from transformers import Qwen3Config
 
 ZEConfig = Qwen3Config
-
-_CE.to = to_device