zzl
commited on
Commit
·
0be014d
1
Parent(s):
44c611e
[Release] cpu support
Browse files- demo_vid.py +1 -1
demo_vid.py
CHANGED
|
@@ -21,7 +21,7 @@ def vid2vid(model_type, video, iters):
|
|
| 21 |
ckpt_path = hf_hub_download(repo_id='lalala125/AMT', filename=f'{model_type.lower()}.pth')
|
| 22 |
print(model_type)
|
| 23 |
ckpt = torch.load(ckpt_path)
|
| 24 |
-
model.load_state_dict(ckpt['state_dict'])
|
| 25 |
model.eval()
|
| 26 |
vcap = cv2.VideoCapture(video)
|
| 27 |
ori_frame_rate = vcap.get(cv2.CAP_PROP_FPS)
|
|
|
|
| 21 |
ckpt_path = hf_hub_download(repo_id='lalala125/AMT', filename=f'{model_type.lower()}.pth')
|
| 22 |
print(model_type)
|
| 23 |
ckpt = torch.load(ckpt_path)
|
| 24 |
+
model.load_state_dict(ckpt['state_dict'], map_location=torch.device('cpu'))
|
| 25 |
model.eval()
|
| 26 |
vcap = cv2.VideoCapture(video)
|
| 27 |
ori_frame_rate = vcap.get(cv2.CAP_PROP_FPS)
|