JiaxinGe
/

Diffusers-BAGEL

Model card Files Files and versions

para-lost commited on Jul 25

Commit

c387723

·

1 Parent(s): 37b1444

fix device

Files changed (1) hide show

pipeline.py +1 -1

pipeline.py CHANGED Viewed

@@ -6819,7 +6819,7 @@ class InterleaveInferencer:
         past_key_values = gen_context['past_key_values']
         kv_lens = gen_context['kv_lens']
         ropes = gen_context['ropes']
         generation_input = self.model.prepare_start_tokens(kv_lens, ropes, self.new_token_ids)
         generation_input = self._to_device(generation_input, device)
         unpacked_latent = self.model.generate_text(

         past_key_values = gen_context['past_key_values']
         kv_lens = gen_context['kv_lens']
         ropes = gen_context['ropes']
+        device = next(self.model.parameters()).device
         generation_input = self.model.prepare_start_tokens(kv_lens, ropes, self.new_token_ids)
         generation_input = self._to_device(generation_input, device)
         unpacked_latent = self.model.generate_text(