update testsingle.py.

2022-10-18 03:59:15 +00:00
parent 9912dca40c
commit 952d7e3a3c
1 changed files with 1 additions and 1 deletions
--- a/testsingle.py
+++ b/testsingle.py
@ -36,7 +36,7 @@ model = None
 #model = VisionTransformer(config, args.img_size, zero_head=True, num_classes=num_classes, smoothing_value=args.smoothing_value)

 if args.pretrained_model is not None:
-    model = torch.load(args.pretrained_model) #自己预训练模型
+    model = torch.load(args.pretrained_model,map_location=torch.device('cpu')) #自己预训练模型
 model.to(args.device)
 model.eval()