Update libs/test_base.py

libs/test_base.py

@@ -15,7 +15,7 @@ class TesterBase():
     def __init__(self, args):
         cudnn.benchmark = True
 
-        self.mean_values = torch.tensor([0.411, 0.432, 0.45]).view(1, 3, 1, 1).cuda()
+        self.mean_values = torch.tensor([0.411, 0.432, 0.45]).view(1, 3, 1, 1).to(args.device)
 
         self.args = args