Spaces:
Runtime error
Runtime error
hcs
commited on
Commit
·
b8a0eab
1
Parent(s):
3f481af
Add application file
Browse files
core.py
CHANGED
|
@@ -13,7 +13,7 @@ def load_model(pretrained_dict, new):
|
|
| 13 |
new.load_state_dict(model_dict)
|
| 14 |
|
| 15 |
|
| 16 |
-
if
|
| 17 |
device = torch.device("cuda")
|
| 18 |
else:
|
| 19 |
device = torch.device("cpu")
|
|
|
|
| 13 |
new.load_state_dict(model_dict)
|
| 14 |
|
| 15 |
|
| 16 |
+
if torch.cuda.is_available():
|
| 17 |
device = torch.device("cuda")
|
| 18 |
else:
|
| 19 |
device = torch.device("cpu")
|