Fix CI

src/levanter/main/viz_logprobs.py

@@ -74,7 +74,7 @@ def main(config: VizGpt2Config):
         def compute_log_probs(model: LmHeadModel, example: LmExample):
             model = inference_mode(model, True)
             model = mp.cast_to_compute(model)
-            logprobs = compute_next_token_loss(model, example, reduction=None)
+            logprobs, where, _ = compute_next_token_loss(model, example)
             # roll forward to get the loss for each predicted token
             logprobs = hax.roll(logprobs, 1, Pos)
             return logprobs.rearrange((EvalBatch, Pos)).array