Locutusque commited on
Commit
cf082dc
1 Parent(s): 56d0515

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -5,10 +5,10 @@ import torch
5
  from threading import Thread
6
  import os
7
 
8
- @spaces.GPU
9
  def load_model(model_name):
10
  return pipeline("text-generation", model=model_name, device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True, token=os.environ["token"], use_fast=True)
11
- @spaces.GPU()
12
  def generate(
13
  message,
14
  history,
 
5
  from threading import Thread
6
  import os
7
 
8
+ @spaces.GPU()
9
  def load_model(model_name):
10
  return pipeline("text-generation", model=model_name, device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True, token=os.environ["token"], use_fast=True)
11
+ @spaces.GPU(duration=120)
12
  def generate(
13
  message,
14
  history,