Set device in torch.arange call for .embed()

curt-tigges · Oct 12, 2023 · 3c12d0a · 3c12d0a
1 parent 49fe73b
commit 3c12d0a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/utils/residual_stream.py b/utils/residual_stream.py
@@ -161,7 +161,9 @@ def embed(
             names_filter = lambda name: hook == name, batch_size=batch_size
         )
         out: Float[Tensor, "batch pos d_model"] = cache[hook]
-        return out[torch.arange(len(out)), self.position, :].detach().cpu()
+        return out[
+            torch.arange(len(out), device=out.device), self.position, :
+        ].detach().cpu()
 
     @classmethod
     def get_dataset(