feat(*): set space gpu

app.py

@@ -5,6 +5,10 @@ import torch
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
 
+testgpu= torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print(testgpu) 
+
 @spaces.GPU
 def greet(n):
     print(zero.device) # <-- 'cuda:0' 🤗

config.py

@@ -31,5 +31,5 @@ class Config(object):
         self.weight_decay = 0.01
         self.max_grad_norm = 1.0
         self.learning_rate = 5e-5
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 

main.py

@@ -1,5 +1,6 @@
 # coding: UTF-8
 
+import spaces
 import os
 import time
 import torch
@@ -27,6 +28,7 @@ def set_seed(seed):
     torch.cuda.manual_seed_all(seed)
     torch.backends.cudnn.deterministic = True
 
+@spaces.GPU
 def main():
     set_seed(args.seed)
     config = Config(args.data_dir)