Update app.py
Browse files
app.py
CHANGED
@@ -43,7 +43,7 @@ yaml=YAML(typ='safe')
|
|
43 |
|
44 |
|
45 |
use_cuda = torch.cuda.is_available()
|
46 |
-
device = torch.
|
47 |
|
48 |
## Load model
|
49 |
|
|
|
43 |
|
44 |
|
45 |
use_cuda = torch.cuda.is_available()
|
46 |
+
device = torch.device('cuda') if use_cuda else torch.device('cpu')
|
47 |
|
48 |
## Load model
|
49 |
|