device bugfix

cybershiptrooper · Sep 7, 2024 · aea84c7 · aea84c7
1 parent e4c4448
commit aea84c7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/iit/tasks/ioi/utils.py b/iit/tasks/ioi/utils.py
@@ -29,7 +29,7 @@ def make_ioi_dataset_and_hl(
 
     ioi_names = t.tensor(
         [ll_model.tokenizer.encode(" " + name) for name in ioi_dataset_tl.names]
-    ).flatten()
+    ).flatten().to(device)
     hl_model = IOI_HL(d_vocab=ll_model.cfg.d_vocab_out, names=ioi_names, device=device)
 
     ioi_dataset = IOIDatasetWrapper(