move model and input ids to GPU

app.py

@@ -110,8 +110,8 @@ def get_stream_warning_error(stream):
 #     return chat_completion
 
 def model_inference(messages):
-    # input_ids = tokenizer(get_last_question(), return_tensors="pt").input_ids.to("cuda") # tokenize the input and put it on the GPU
-    input_ids = tokenizer(get_last_question(), return_tensors="pt").input_ids # move to GPU eventually
+    input_ids = tokenizer(get_last_question(), return_tensors="pt").input_ids.to("cuda") # tokenize the input and put it on the GPU
+    # input_ids = tokenizer(get_last_question(), return_tensors="pt").input_ids # testing on CPU
     outputs = model.generate(input_ids)
     for chunk in tokenizer.decode(outputs[0], skip_special_tokens=True):
         yield chunk # yield each chunk of the predicted string character by character

model.py

@@ -27,8 +27,8 @@ class InferenceBuilder:
         # cannot directly use @st.cache_resource on a method (function within a class) that has a self argument
         @st.cache_resource # https://docs.streamlit.io/develop/concepts/architecture/caching
         def load_and_cache_model(model_name):
-            # model = T5ForConditionalGeneration.from_pretrained(model_name).to("cuda") # put the model on our Space's GPU
-            model = T5ForConditionalGeneration.from_pretrained(model_name) # move to GPU eventually
+            model = T5ForConditionalGeneration.from_pretrained(model_name).to("cuda") # put the model on our Space's GPU
+            # model = T5ForConditionalGeneration.from_pretrained(model_name) # testing on CPU
             return model
         
         return load_and_cache_model(model_name)