Put specifying args.device in main file instead of inside function

app.py

@@ -103,10 +103,6 @@ def get_args_parser():
 @spaces.GPU
 # Get counting model.
 def build_model_and_transforms(args):
-    if torch.cuda.is_available():
-        args.device = torch.device('cpu')
-    else:
-        args.device = torch.device('cpu')
     normalize = T.Compose(
         [T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]
     )
@@ -153,6 +149,10 @@ def build_model_and_transforms(args):
 
 parser = argparse.ArgumentParser("Counting Application", parents=[get_args_parser()])
 args = parser.parse_args()
+if torch.cuda.is_available():
+    args.device = torch.device('cpu')
+else:
+    args.device = torch.device('cpu')
 model, transform = build_model_and_transforms(args)
 
 examples = [