Update app.py

app.py

@@ -27,7 +27,7 @@ def inference(context, instruction, num_beams:int=4):
     if use_cuda:
         for t in input_tokens:
             if torch.is_tensor(input_tokens[t]):
-                input_tokens[t] = input_tokens[t].to(input_device)
+                input_tokens[t] = input_tokens[t].to("cuda")
 
     output = model.generate(
         input_tokens['input_ids'],
@@ -42,7 +42,7 @@ def inference(context, instruction, num_beams:int=4):
     output_tokens = output.sequences
     generated_tokens = output_tokens[:, num_input_tokens:]
     num_generated_tokens = (generated_tokens != tokenizer.pad_token_id).sum(dim=-1).tolist()[0]
-    prefix_to_add = torch.tensor([[tokenizer("A")["input_ids"][0]]]).to(input_device)
+    prefix_to_add = torch.tensor([[tokenizer("A")["input_ids"][0]]]).to("cuda")
     generated_tokens = torch.cat([prefix_to_add, generated_tokens], dim=1)
     generated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
     string_output = [i[1:].strip() for i in generated_text][0]