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")