Take cache mode activations off the GPU (#355)

ApolloResearch · Mar 11, 2024 · a3f4003 · a3f4003
1 parent 68a2fc7
commit a3f4003
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rib/hook_fns.py b/rib/hook_fns.py
@@ -66,7 +66,7 @@ def _append_to_hooked_list(
         element_to_append: Appended to hooked data.
     """
     hooked_data.setdefault(hook_name, {}).setdefault(data_key, [])
-    hooked_data[hook_name][data_key].append(element_to_append)
+    hooked_data[hook_name][data_key].append(element_to_append.cpu())
 
 
 InputActType = Union[