Update interpret.py

interpret.py

@@ -92,7 +92,7 @@ class InterpretationPrompt:
 
     def generate(self, model, embeds, k, layer_format='model.layers.{k}', **generation_kwargs):
         num_seqs = len(embeds[0]) # assumes the placeholder 0 exists
-        tokens_batch = torch.tensor([self.tokens[:] for _ in range(num_seqs)])
+        tokens_batch = torch.tensor([self.tokens[:] for _ in range(num_seqs)]).to(model.device)
         module = model.get_submodule(layer_format.format(k=k))
         with SubstitutionHook(module, positions_dict=self.placeholders, values_dict=embeds):
             generated = model.generate(tokens_batch, **generation_kwargs)