fix typing error

TransformerLensOrg · Sep 26, 2024 · ad0dea0 · ad0dea0
1 parent e618d44
commit ad0dea0
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/transformer_lens/utils.py b/transformer_lens/utils.py
@@ -820,8 +820,10 @@ def test_prompt(
 
  # If n_answers = 1 then unwrap answer ranks, so printed output matches original version of function
  if not using_multiple_answers:
- answer_ranks = [r[0] for r in answer_ranks]
- rprint(f"[b]Ranks of the answer tokens:[/b] {answer_ranks}")
+ single_answer_ranks = [r[0] for r in answer_ranks]
+ rprint(f"[b]Ranks of the answer tokens:[/b] {single_answer_ranks}")
+ else:
+ rprint(f"[b]Ranks of the answer tokens:[/b] {answer_ranks}")
 
 
 def transpose(tensor: Float[torch.Tensor, "... a b"]) -> Float[torch.Tensor, "... b a"]: