Update app.py

app.py

@@ -69,6 +69,9 @@ def generate_prompt(instruction, input=None):
         return PROMPT_DICT["prompt_no_input"].format(instruction=instruction)
 
 @spaces.GPU
+def getIds(inputs):
+    return inputs["input_ids"].cuda()
+
 def generator(input_ids, generation_config, max_new_tokens):
     # Without streaming
     with torch.no_grad():
@@ -91,7 +94,7 @@ def respond(
 ):
     ins_f = generate_prompt(message,None)
     inputs  =  tokenizer(ins_f, return_tensors="pt")
-    input_ids = inputs["input_ids"].cuda()
+    input_ids = getIds(inputs)
     max_new_tokens = 512
     generation_config = GenerationConfig(
             temperature=0.1,