Updated to remove hardcoded .cpu() processing#1189

jlarson4 merged 1 commit intodev-3.x-canaryTransformerLensOrg/TransformerLens:dev-3.x-canaryfrom

bug/cache-device-run-with-cacheTransformerLensOrg/TransformerLens:bug/cache-device-run-with-cacheCopy head branch name to clipboard