YuAnthony commited on
Commit
6e41f12
·
1 Parent(s): eb3c3fe

Update infer_contrast.py

Browse files
Files changed (1) hide show
  1. infer_contrast.py +2 -2
infer_contrast.py CHANGED
@@ -17,11 +17,11 @@ args =parser.parse_known_args()[0]
17
 
18
  print_arguments(args)
19
 
20
- device = torch.device("cuda")
21
 
22
  # 加载模型
23
  # model = torch.jit.load(args.model_path)
24
- model = torch.load(args.model_path)
25
  # model.to(device)
26
  model.eval()
27
 
 
17
 
18
  print_arguments(args)
19
 
20
+ device = torch.device("cpu")
21
 
22
  # 加载模型
23
  # model = torch.jit.load(args.model_path)
24
+ model = torch.load(args.model_path,device=device)
25
  # model.to(device)
26
  model.eval()
27