diff --git a/mu_map/training/random_search.py b/mu_map/training/random_search.py
index e6b57b3bcfdaf438c08d6fec6e8141dd562beb12..5f87abadd00bf7237ddcc4b3072b2f7566ff8eab 100644
--- a/mu_map/training/random_search.py
+++ b/mu_map/training/random_search.py
@@ -239,8 +239,7 @@ class RandomSearchCGAN(RandomSearch):
         self.iterations = iterations
         self.dir = "cgan_random_search"
         validate_and_make_directory(self.dir)
-        # self.device = torch.device("cuda")
-        self.device = torch.device("cpu")
+        self.device = torch.device("cuda")
         self.params = {}
 
         self.dir_train = os.path.join(self.dir, "train_data")