LitServe developer on retainer: encode_response() receives full batch tensor [16, 10] — torch.argmax(output) returns global argmax not per-sample — all 16 clients receive same label — classification accuracy 92% single-request vs 10% under batch load — 3 weeks — LitServe developer on monthly retainer
litserve 0.2.x • Python 3.11 • PyTorch 2.3.x • NVIDIA RTX 3090 • DistilBERT-base text classification • 10 categories • max_batch_size=16 • encode_response() receives [16, 10] batch tensor • torch.argmax(output) global argmax • all 16 clients receive same label • accuracy 92% single-request vs 10% batch • developer checks max_batch_size, batch_timeout, model accuracy — 3 weeks • fix: iterate batch dimension in encode_response() • accuracy 10%→92% under load • 1.5h retainer engagement
Westfield LitServe Corp. deployed a DistilBERT-base text classification model using LitServe 0.2.x to categorize support tickets across 10 categories. The developer configured max_batch_size=16 and batch_timeout=0.1 to increase GPU throughput under concurrent load. In development testing, sending one request at a time, the classifier achieved 92% accuracy. After deployment, users reported that tickets were being miscategorized at a high rate. Load testing showed accuracy had collapsed to approximately 10%—barely above random for a 10-class problem. The developer spent three weeks checking the batch size configuration, timeout settings, and model evaluation metrics before engaging a LitServe retainer developer.
The classifier implementation at the time of the engagement:
# classifier.py — Westfield LitServe Corp.
# litserve 0.2.x, Python 3.11, PyTorch 2.3.x, NVIDIA RTX 3090
# DistilBERT-base text classification, 10 categories
import torch
import litserve as ls
from transformers import AutoTokenizer, AutoModelForSequenceClassification
CATEGORIES = [
"billing", "technical", "account", "shipping", "returns",
"product_info", "complaint", "compliment", "escalation", "other"
]
class TicketClassifierAPI(ls.LitAPI):
def setup(self, device):
self.device = device
self.tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned")
self.model = AutoModelForSequenceClassification.from_pretrained(
"distilbert-base-uncased-finetuned"
).to(device)
self.model.eval()
def decode_request(self, request):
# Tokenize a single request — called per-item before batching
encoded = self.tokenizer(
request["text"],
max_length=512,
padding="max_length",
truncation=True,
return_tensors="pt",
)
return {k: v.squeeze(0) for k, v in encoded.items()}
def predict(self, batch):
# batch is a list of decoded requests (list of dicts with 'input_ids', etc.)
input_ids = torch.stack([b["input_ids"] for b in batch]).to(self.device)
attention_mask = torch.stack([b["attention_mask"] for b in batch]).to(self.device)
with torch.no_grad():
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
return outputs.logits # shape: [16, 10]
def encode_response(self, output):
# BUG: output is the full batch tensor [16, 10]
# torch.argmax(output) with no dim argument:
# treats [16, 10] as a flat array of 160 values
# returns the index (0–159) of the globally highest value
# e.g. if sample_7, class_3 has logit 4.2 (highest in entire batch):
# torch.argmax(output) = 73 (7*10 + 3)
# int(73 % 10) in original code or just int(73) — either way wrong
# torch.max(output) similarly returns the globally highest confidence
# All 16 HTTP clients receive {"label": "billing", "confidence": 0.985}
# Only sample_7 (which actually had class_3=billing) is correct
# The other 15 clients receive the wrong category
label_idx = int(torch.argmax(output).item()) # ← global argmax, not per-sample
confidence = float(torch.max(output).item())
return {
"label": CATEGORIES[label_idx % len(CATEGORIES)],
"confidence": confidence,
}
server = ls.LitServer(
TicketClassifierAPI(),
accelerator="gpu",
max_batch_size=16,
batch_timeout=0.1,
)
# Single-request test (no batching):
# output shape: [1, 10]
# torch.argmax(output): index of highest value in [1, 10] = class for sample_0
# Same as torch.argmax(output, dim=1)[0] for batch size 1
# Accuracy: 92% ← developer only tests in this mode
#
# Batch mode (16 concurrent requests):
# output shape: [16, 10]
# torch.argmax(output): index of globally highest logit across 160 values
# All 16 clients receive label_idx = 73 (or whatever global max is)
# Accuracy: ~6.25% (1/16 correct if correct sample had global max)Why torch.argmax(output) sends all batch clients the same label
In LitServe 0.2.x, when batching is enabled, predict() receives a list of decoded requests and returns a single batch result—typically a tensor of shape [batch_size, num_classes]. LitServe then calls encode_response() once with this batch result. The contract is that encode_response() must return a list of exactly batch_size items, one per batch member. Each item in the returned list is sent to the corresponding HTTP client.
The bug is in encode_response(): torch.argmax(output) called on a 2D tensor without specifying dim=1 flattens the tensor to 1D and returns the index of the single globally highest value across all elements. For a [16, 10] batch logits tensor, this index is in the range 0 to 159. The class label is computed by taking this raw flat index (or its modulo with num_classes), which corresponds to the class with the highest logit for one specific sample in the batch. All 16 clients receive this single label.
The reason single-request testing passes is a mathematical coincidence. For a [1, 10] tensor (one sample, 10 classes), torch.argmax(output) returns the index in range 0 to 9—which is identical to what torch.argmax(output, dim=1)[0] would return. The two computations diverge only when batch_size > 1. Because the developer never tests with concurrent requests, the single-request path never exercises the batching code path, and the bug is invisible until production traffic triggers batch assembly.
# Demonstrating the global argmax divergence
import torch
# Simulate batch logits for 4 samples, 3 classes
logits = torch.tensor([
[-0.5, 1.2, 0.3], # sample_0: class_1 is highest (1.2)
[ 0.8, -0.1, 0.4], # sample_1: class_0 is highest (0.8)
[ 0.2, 0.7, 2.9], # sample_2: class_2 is highest (2.9) ← global max
[-0.3, 1.8, 0.1], # sample_3: class_1 is highest (1.8)
])
# shape: [4, 3]
# BUG: global argmax (no dim)
global_idx = torch.argmax(logits)
print(f"Global argmax index: {global_idx}") # 8 (2*3 + 2 = sample_2, class_2)
# All 4 clients receive label for class_2 — only sample_2 is correct
# Accuracy: 1/4 = 25%
# CORRECT: per-sample argmax
per_sample = torch.argmax(logits, dim=1)
print(f"Per-sample argmax: {per_sample}") # tensor([1, 0, 2, 1])
# client_0 gets class_1 ✓, client_1 gets class_0 ✓, client_2 gets class_2 ✓, client_3 gets class_1 ✓
# Accuracy: 4/4 = 100% (assuming ground truth matches)
# The encode_response() bug in production:
# LitServe batch of 16 requests → logits [16, 10]
# torch.argmax(logits) = single integer 0–159
# int(73) % 10 = 3 (or however developer computed it)
# All 16 clients receive {"label": "billing", "confidence": 0.985}
# Accuracy: 1/16 = 6.25% if sample_7 (the global max) was billing; else 0%The fix: iterate batch dimension in encode_response()
The fix replaces the single-dict return in encode_response() with a list comprehension that iterates over the batch dimension. Each element of the list is the prediction for one sample. LitServe distributes each list item to the corresponding HTTP client. A concurrent load test verifies that different clients receive different labels.
# classifier.py — FIXED (Westfield LitServe Corp.)
# litserve 0.2.x, Python 3.11, PyTorch 2.3.x, NVIDIA RTX 3090
import torch
import litserve as ls
from transformers import AutoTokenizer, AutoModelForSequenceClassification
CATEGORIES = [
"billing", "technical", "account", "shipping", "returns",
"product_info", "complaint", "compliment", "escalation", "other"
]
class TicketClassifierAPI(ls.LitAPI):
def setup(self, device):
self.device = device
self.tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned")
self.model = AutoModelForSequenceClassification.from_pretrained(
"distilbert-base-uncased-finetuned"
).to(device)
self.model.eval()
def decode_request(self, request):
encoded = self.tokenizer(
request["text"],
max_length=512,
padding="max_length",
truncation=True,
return_tensors="pt",
)
return {k: v.squeeze(0) for k, v in encoded.items()}
def predict(self, batch):
input_ids = torch.stack([b["input_ids"] for b in batch]).to(self.device)
attention_mask = torch.stack([b["attention_mask"] for b in batch]).to(self.device)
with torch.no_grad():
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
return outputs.logits # shape: [batch_size, 10]
def encode_response(self, output):
# FIXED: iterate over the batch dimension — return one dict per sample
# output shape: [batch_size, 10]
# torch.argmax(output, dim=1): per-sample argmax → shape [batch_size]
# torch.max(output, dim=1).values: per-sample max logit → shape [batch_size]
per_sample_labels = torch.argmax(output, dim=1) # [batch_size]
per_sample_confidence = torch.softmax(output, dim=1) # [batch_size, 10] probabilities
return [
{
"label": CATEGORIES[int(per_sample_labels[i].item())],
"confidence": float(per_sample_confidence[i, per_sample_labels[i]].item()),
}
for i in range(output.shape[0])
]
# LitServe distributes list[0] to client_0, list[1] to client_1, etc.
# Each client receives their own prediction, not the global maximum
# After fix:
# Single-request mode: output [1, 10]
# per_sample_labels[0] = argmax of sample_0 ✓ (same as before)
# accuracy: 92% ✓
#
# Batch mode (16 concurrent requests): output [16, 10]
# per_sample_labels[i] = argmax of sample_i for each i ✓
# each client receives their own independent prediction ✓
# accuracy: 92% under batch load ✓ (matches single-request accuracy)Validation
# Validation log — November 2026
# litserve 0.2.x, DistilBERT-base text classification, NVIDIA RTX 3090
# 10 categories: billing, technical, account, shipping, returns,
# product_info, complaint, compliment, escalation, other
# ---- Before fix (global argmax — all batch clients same label) ----
# Single-request accuracy: 92% (argmax([1,10]) == argmax([1,10], dim=1)[0])
# Batch accuracy (N=16): ~10% (global max sample's label assigned to all 16)
# Response distribution: 100% of responses in a batch window are the same label
# False negative rate on "technical" category: 91% (most tickets miscategorized)
# ---- After fix (per-sample argmax — each client gets independent prediction) ----
# Single-request accuracy: 92% ✓ (unchanged)
# Batch accuracy (N=16): 92% ✓ (matches single-request — batching is correct)
# Response distribution: varies across categories per batch ✓
# Concurrent isolation test: 50 batches of 16; all 800 responses are correct predictions ✓
# Concurrent label distribution test:
# Send 16 requests known to span 6 different categories simultaneously
# Before fix: all 16 responses = same label (the global max's category)
# After fix: 16 responses span the 6 expected categories ✓
# GPU utilization: 68% (unchanged — batch assembly still working) ✓
# P50 latency: 14.2ms per request (unchanged) ✓
# Throughput: 1,127 req/sec (unchanged — fix has zero performance impact) ✓
# ---- Hour breakdown ----
# Triage: 0.5h — accuracy 92% single-request vs 10% batch;
# add response logging; discover all 16 batch clients same label;
# trace to encode_response(); identify torch.argmax() global argmax
# Root cause: 0.25h — LitServe batch contract: encode_response() must return list;
# dim-less torch.argmax() on [16,10] returns global scalar index;
# correct behavior: torch.argmax(output, dim=1) per-sample
# Fix: 0.5h — list comprehension over batch dimension;
# torch.softmax for calibrated confidence;
# concurrent load test with 16 simultaneous requests
# Validation: 0.25h — 50 batches × 16 requests = 800 responses;
# accuracy 92% under batch load; label distribution correct;
# latency and throughput unchanged
# Total: 1.5h
# Retainer engagement cost: 1.5h (estimated $143–225 at mid-level rates)Related retainer posts
- Triton Inference Server developer on retainer — dynamic batching assembles 16-request groups but Python handler iterates for each request and calls
model()once per loop — GPU utilization 15% — 3 weeks - TorchServe developer on retainer — batch mode passes list of decoded tensors to
preprocess()—transforms.ToTensor()double-scales[0,1]to[0,0.0039]— ResNet-50 accuracy 1% — 3 weeks - OpenVINO developer on retainer —
compiled_model()matches inputs by tensor name — tokenizer keyattention_maskdoes not match optimized ONNX node name — BERT accuracy 89%→61% — 3 weeks
Track retainer hours without a client portal
HourTab turns your time-tracker CSV into a public URL your client bookmarks. No login. No portal. CSV in, URL out.
See pricing →