Benson commited on
Commit
b60c56b
·
1 Parent(s): f74b786

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -7,7 +7,7 @@ app = FastAPI(docs_url="/")
7
  # Initialize the text generation pipeline
8
  # This function will be able to generate text
9
  # given an input.
10
- pipe = pipeline("text2text-generation", model="google/flan-t5-small")
11
 
12
  # Define a function to handle the GET request at `/generate`
13
  # The generate() function is defined as a FastAPI route that takes a
 
7
  # Initialize the text generation pipeline
8
  # This function will be able to generate text
9
  # given an input.
10
+ pipe = pipeline("text2text-generation", model="google/flan-t5-small",device_map="auto")
11
 
12
  # Define a function to handle the GET request at `/generate`
13
  # The generate() function is defined as a FastAPI route that takes a