[Bug Report] evals.py ablates all heads when it needs to ablate only one

[shehper]

If you are submitting a bug report, please fill in the following details and use the tag [bug].

Describe the bug
When evaluating SAEs trained on individual attention head outputs, the reconstructed score should be calculated by zero-ablating only that singular head.

Here, zero_abl_loss calls zero_ablate_hook, which zero-ablates all heads.

Additional context
I noticed this while training SAEs on individual attention head outputs. The reconstructed score looked really good (~0.97) but the feature dashboards looked quite mediocre.

Checklist

• I have checked that there is no similar issue in the repo (required)

[shehper]

I will be happy to submit a PR for this.

[jbloomAus]

Hey Shepher,

Thanks for raising this. So obvious in hindsight. I'd love a PR for this.

[shehper]

Sorry, I have just been a little swamped with the research sprint. I'll submit a PR over the weekend.