fix default cuda

autogl/module/nas/algorithm/darts.py

@@ -103,7 +103,7 @@ class Darts(BaseNAS):
         model_wd=5e-4,
         arch_lr=3e-4,
         arch_wd=1e-3,
-        device="cuda",
+        device="auto",
     ):
         super().__init__(device=device)
         self.num_epochs = num_epochs

autogl/module/nas/algorithm/enas.py

@@ -79,7 +79,7 @@ class Enas(BaseNAS):
         model_lr=5e-3,
         model_wd=5e-4,
         disable_progress=True,
-        device="cuda",
+        device="auto",
     ):
         super().__init__(device)
         self.device = device

autogl/module/nas/algorithm/gasso.py

@@ -34,7 +34,7 @@ class Gasso(BaseNAS):
         arch_lr = 0.03,
         stru_lr = 0.04,
         lamb = 0.6,
-        device="cuda",
+        device="auto",
     ):
         super().__init__(device=device)
         self.device = device

autogl/module/nas/algorithm/random_search.py

@@ -35,7 +35,7 @@ class RandomSearch(BaseNAS):
         Control whether show the progress bar.
     """
 
-    def __init__(self, device="cuda", num_epochs=400, disable_progress=False, hardware_metric_limit=None):
+    def __init__(self, device="auto", num_epochs=400, disable_progress=False, hardware_metric_limit=None):
         super().__init__(device)
         self.num_epochs = num_epochs
         self.disable_progress = disable_progress

autogl/module/nas/algorithm/rl.py

@@ -250,7 +250,7 @@ class RL(BaseNAS):
     def __init__(
         self,
         num_epochs=5,
-        device="cuda",
+        device="auto",
         log_frequency=None,
         grad_clip=5.0,
         entropy_weight=0.0001,
@@ -429,7 +429,7 @@ class GraphNasRL(BaseNAS):
 
     def __init__(
         self,
-        device="cuda",
+        device="auto",
         num_epochs=10,
         log_frequency=None,
         grad_clip=5.0,

autogl/module/nas/space/base.py

@@ -98,7 +98,7 @@ class BoxModel(BaseModel):
 
     _logger = get_logger("space model")
 
-    def __init__(self, space_model, device=torch.device("cuda")):
+    def __init__(self, space_model, device=torch.device("auto")):
         super().__init__(init=True)
         self.init = True
         self.space = []
@@ -247,7 +247,7 @@ class BaseSpace(nn.Module):
         )
         return layer
 
-    def wrap(self, device="cuda"):
+    def wrap(self, device="auto"):
         """
         Return a BoxModel which wrap self as a model
         Used to pass to trainer

autogl/module/nas/space/gasso.py

@@ -256,12 +256,17 @@ class GassoSpace(BaseSpace):
     def keep_prediction(self):
         self.prediction = self.current_pred
 
+    def to(self, *args, **kwargs):
+        super().to(args, kwargs)
+        device = next(self.parameters()).device
+        self.alphas_normal = self.alphas_normal.to(device)
+
     def initialize_alphas(self):
         num_ops = len(self.ops)
 
         self.alphas_normal = []
         for i in range(self.steps):
-            self.alphas_normal.append(Variable(1e-3 * torch.randn(num_ops).cuda(), requires_grad=True))
+            self.alphas_normal.append(Variable(1e-3 * torch.randn(num_ops), requires_grad=True))
 
         self._arch_parameters = [
             self.alphas_normal