From e495b6f7aef3e897681f6cb80ce2b1f4e5a7a57c Mon Sep 17 00:00:00 2001
From: generall <general502570@outlook.com>
Date: Sat, 25 Dec 2021 11:42:23 +0800
Subject: [PATCH] fix default cuda

---
 autogl/module/nas/algorithm/darts.py         | 2 +-
 autogl/module/nas/algorithm/enas.py          | 2 +-
 autogl/module/nas/algorithm/gasso.py         | 2 +-
 autogl/module/nas/algorithm/random_search.py | 2 +-
 autogl/module/nas/algorithm/rl.py            | 4 ++--
 autogl/module/nas/space/base.py              | 4 ++--
 autogl/module/nas/space/gasso.py             | 7 ++++++-
 7 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/autogl/module/nas/algorithm/darts.py b/autogl/module/nas/algorithm/darts.py
index 87627e7..4d85292 100644
--- a/autogl/module/nas/algorithm/darts.py
+++ b/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
diff --git a/autogl/module/nas/algorithm/enas.py b/autogl/module/nas/algorithm/enas.py
index 36fac6e..d2c8dc3 100644
--- a/autogl/module/nas/algorithm/enas.py
+++ b/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
diff --git a/autogl/module/nas/algorithm/gasso.py b/autogl/module/nas/algorithm/gasso.py
index 42345cc..fc50ed6 100644
--- a/autogl/module/nas/algorithm/gasso.py
+++ b/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
diff --git a/autogl/module/nas/algorithm/random_search.py b/autogl/module/nas/algorithm/random_search.py
index f704dbd..8abd292 100644
--- a/autogl/module/nas/algorithm/random_search.py
+++ b/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
diff --git a/autogl/module/nas/algorithm/rl.py b/autogl/module/nas/algorithm/rl.py
index 2f8e9a8..d927c26 100644
--- a/autogl/module/nas/algorithm/rl.py
+++ b/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,
diff --git a/autogl/module/nas/space/base.py b/autogl/module/nas/space/base.py
index 8ff7901..e122ccb 100644
--- a/autogl/module/nas/space/base.py
+++ b/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
diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py
index 0397cb5..9d4db84 100644
--- a/autogl/module/nas/space/gasso.py
+++ b/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