Update RIFE.py

2022-11-16 11:59:13 +08:00 · 2022-11-16 11:59:13 +08:00 · d430234d81
parent 1cc683d8e2
commit d430234d81
1 changed files with 1 additions and 1 deletions
--- a/model/RIFE.py
+++ b/model/RIFE.py
@ -13,7 +13,7 @@ from model.loss import *
 from model.laplacian import *
 from model.refine import *

-device = torch.device("cuda")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
 class Model:
    def __init__(self, local_rank=-1, arbitrary=False):