fix: detach grads from gated SAE for aux loss

[chanind]

Description

This PR detaches gradients from the gated sparsity loss in gated SAEs. Thanks to @jasonrichdarmawan for finding this!

Fixes #539

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

You have tested formatting, typing and tests

• I have run make check-ci to check format and linting. (you can run make format to format code if needed.)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.83%. Comparing base (3432f00) to head (a231b21).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #545   +/-   ##
=======================================
  Coverage   84.83%   84.83%           
=======================================
  Files          29       29           
  Lines        3766     3766           
  Branches      448      448           
=======================================
  Hits         3195     3195           
  Misses        397      397           
  Partials      174      174           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.