philschmid HF staff commited on
Commit
4f89bf3
1 Parent(s): 57e9c15

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -4,7 +4,7 @@ import torch
4
 
5
 
6
  # check for GPU
7
- device = "cuda" if torch.cuda.is_available() else "cpu"
8
 
9
 
10
  class EndpointHandler:
 
4
 
5
 
6
  # check for GPU
7
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
 
9
 
10
  class EndpointHandler: