Use inference_mode as decorator

app.py

@@ -67,6 +67,7 @@ inputs = tokenizer([text])
 window_len = min(window_len, len(input_ids))
 
 @st.cache_data(show_spinner=False)
+@torch.inference_mode()
 def run_context_length_probing(model_name, text, window_len):
     assert model.name_or_path == model_name
     del text  # needed as a cache key but for the computation we access inputs directly
@@ -76,12 +77,13 @@ def run_context_length_probing(model_name, text, window_len):
         window_len=window_len,
         pad_id=tokenizer.eos_token_id
     )
-    with torch.inference_mode():
-        logits = model(**inputs_sliding.convert_to_tensors("pt")).logits.to(torch.float16)
-        logits = logits.permute(1, 0, 2)
-        logits = F.pad(logits, (0, 0, 0, window_len, 0, 0), value=torch.nan)
-        logits = logits.view(-1, logits.shape[-1])[:-window_len]
-        logits = logits.view(window_len, len(input_ids) + window_len - 2, logits.shape[-1])
+
+    logits = model(**inputs_sliding.convert_to_tensors("pt")).logits.to(torch.float16)
+
+    logits = logits.permute(1, 0, 2)
+    logits = F.pad(logits, (0, 0, 0, window_len, 0, 0), value=torch.nan)
+    logits = logits.view(-1, logits.shape[-1])[:-window_len]
+    logits = logits.view(window_len, len(input_ids) + window_len - 2, logits.shape[-1])
 
     scores = logits.to(torch.float32).log_softmax(dim=-1)
     scores = scores[:, torch.arange(len(input_ids[1:])), input_ids[1:]]