Put prompt_token_ids, attentions_mask and weights on the same d…

…evice

outlines/generate/api.py

@@ -279,13 +279,9 @@ def stream(
  stop_sequences = stop_at
  num_samples = self.num_samples
 
- if rng is None:
- rng = torch.Generator(device=self.device)
- rng.seed()
-
  prompt_token_ids, attention_masks = self.tokenizer.encode(prompts)
  prompt_token_ids = prompt_token_ids.to(self.device)
- attention_masks = attention_masks.to(self.device)
+ attention_masks = attention_masks.to(prompt_token_ids.device)
 
  # To draw multiple samples we repeat the prompt as many times
  # as there are samples. We copy the FSMs and initialize the
@@ -298,9 +294,15 @@ def stream(
  fsm_states = [FSMState(0) for _ in range(batch_size * num_samples)]
  fsms = [self.fsm.copy() for _ in range(batch_size * num_samples)]
  weights = torch.zeros(
- (batch_size * num_samples), dtype=torch.float, device=self.device
+ (batch_size * num_samples),
+ dtype=torch.float,
+ device=prompt_token_ids.device,
  )
 
+ if rng is None:
+ rng = torch.Generator(device=prompt_token_ids.device)
+ rng.seed()
+
  states = sequence_generator(
  self.model,
  self.sampler,