Spaces:

chenjoya
/

LiveCC

Running on Zero

chenjoya commited on Apr 23

Commit

821526b

verified ·

1 Parent(s): fddefb0

Update demo/infer.py

Files changed (1) hide show

demo/infer.py CHANGED Viewed

@@ -206,7 +206,7 @@ class LiveCCDemoInfer:
             return_tensors="pt",
             return_attention_mask=False
         )
-        inputs.to('cuda')
         if past_ids is not None:
             inputs['input_ids'] = torch.cat([past_ids, inputs.input_ids], dim=1)
         outputs = self.model.generate(

             return_tensors="pt",
             return_attention_mask=False
         )
+        inputs.to(self.model.device)
         if past_ids is not None:
             inputs['input_ids'] = torch.cat([past_ids, inputs.input_ids], dim=1)
         outputs = self.model.generate(