Spaces:

sunshineatnoon
/

TextureScraping

Runtime error

sunshineatnoon commited on Apr 24, 2022

Commit

8fb0c3a

1 Parent(s): bfbfc21

Update libs/test_base.py

Files changed (1) hide show

libs/test_base.py CHANGED Viewed

@@ -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.args = args

     def __init__(self, args):
         cudnn.benchmark = True
+        self.mean_values = torch.tensor([0.411, 0.432, 0.45]).view(1, 3, 1, 1).to(args.device)
         self.args = args