Fix TPU evaluation inference-mode crash

[charlesmartin14]

Summary

• replace torch.inference_mode() with torch.no_grad() in evaluate_probe
• apply the same compatibility fix to evaluate_bleu
• preserve no-autograd evaluation semantics while avoiding PyTorch/XLA tensor-subclass inference-tensor version-counter failures

Reproduction

On a TPU v5e (v5litepod-4) with torch 2.6.0 / torch_xla 2.6.0, the quick smoke reached its first evaluation forward pass and failed with:

RuntimeError: Cannot set version_counter for inference tensor

The stack terminates in evaluate_probe -> model -> q_proj under the @torch.inference_mode() decorator.

Rationale

PyTorch inference mode has known incompatibilities with tensor subclasses. XLA tensors are tensor subclasses. torch.no_grad() disables autograd without converting intermediate tensors to inference tensors, which is appropriate for this metric path.