TastyRice commited on
Commit
0caf86f
1 Parent(s): 5bff417

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -36,7 +36,8 @@ MODEL = os.getenv(
36
  torch.backends.cudnn.deterministic = True
37
  torch.backends.cudnn.benchmark = False
38
 
39
- device = torch.device("cuda:0" if torch.cuda.is_available() else torch.cpu.device.is_available())
 
40
 
41
 
42
  def load_pipeline(model_name):
 
36
  torch.backends.cudnn.deterministic = True
37
  torch.backends.cudnn.benchmark = False
38
 
39
+ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
40
+ model = model.to(device)
41
 
42
 
43
  def load_pipeline(model_name):