Spaces:
Sleeping
Sleeping
Commit ·
e10b61a
1
Parent(s): 25981df
Change to cpu
Browse files- extract_video.py +1 -1
extract_video.py
CHANGED
|
@@ -226,7 +226,7 @@ if __name__ == '__main__':
|
|
| 226 |
|
| 227 |
# net and model
|
| 228 |
net = RetinaFace(cfg=cfg, phase='test')
|
| 229 |
-
net = load_model(net, pretrained_weights, load_to_cpu=False)
|
| 230 |
net.eval()
|
| 231 |
print('Finished loading model!')
|
| 232 |
|
|
|
|
| 226 |
|
| 227 |
# net and model
|
| 228 |
net = RetinaFace(cfg=cfg, phase='test')
|
| 229 |
+
net = load_model(net, pretrained_weights, load_to_cpu=False if device=='cuda' else True)
|
| 230 |
net.eval()
|
| 231 |
print('Finished loading model!')
|
| 232 |
|