Skip to content

Updated to remove hardcoded .cpu() processing#1189

Merged
jlarson4 merged 1 commit intodev-3.x-canaryfrom
bug/cache-device-run-with-cache
Mar 3, 2026
Merged

Updated to remove hardcoded .cpu() processing#1189
jlarson4 merged 1 commit intodev-3.x-canaryfrom
bug/cache-device-run-with-cache

Conversation

@jlarson4
Copy link
Collaborator

@jlarson4 jlarson4 commented Mar 3, 2026

Description

TransformerBridge's run_with_cache was always storing the cache in CPU. In order to maintain consistency with HookedTransformer behavior, updated make_cache_hook to reference the target device.

Fixes # (issue)

Type of change

  • Bug fix (non-breaking change which fixes an issue)

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 unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4 jlarson4 merged commit a212f71 into dev-3.x-canary Mar 3, 2026
29 of 30 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant