Update app.py

app.py

@@ -1,13 +1,17 @@
 from fastapi import FastAPI
 from transformers import pipeline
+import torch
 
 # Create a new FastAPI app instance
 app = FastAPI(docs_url="/")
 
+gpu = torch.device("cuda")
+
+
 # Initialize the text generation pipeline
 # This function will be able to generate text
 # given an input.
-pipe = pipeline("text2text-generation", model="google/flan-t5-large",device=0)
+pipe = pipeline("text2text-generation", model="google/flan-t5-small",device=gpu)
 
 # Define a function to handle the GET request at `/generate`
 # The generate() function is defined as a FastAPI route that takes a 
@@ -25,4 +29,4 @@ def generate(text: str):
     output = pipe(text)
     print(output)
     # Return the generated text in a JSON response
-    return {"output": output[0]["generated_text"]}
+    return {"output": output[0]["generated_text"]}