diff --git a/makefile b/makefile
index e290d11..09b9dbe 100644
--- a/makefile
+++ b/makefile
@@ -17,4 +17,5 @@ animate:
 
 clean:
 	rm -rf lightning_logs/*
-	rm out/*.png
+	rm -f out/*.png
+	cp hsv.png out/
diff --git a/model.py b/model.py
index fd575e2..68b9d7b 100644
--- a/model.py
+++ b/model.py
@@ -86,15 +86,9 @@ class ColorTransformerModel(pl.LightningModule):
 
         # Neural network layers
         self.network = nn.Sequential(
-            nn.Linear(3, 16),
+            nn.Linear(3, 256),
             nn.ReLU(),
-            nn.Linear(16, 16),
-            nn.ReLU(),
-            nn.Linear(16, 128),
-            nn.ReLU(),
-            nn.Linear(128, 128),
-            nn.ReLU(),
-            nn.Linear(128, 64),
+            nn.Linear(256, 64),
             nn.ReLU(),
             nn.Linear(64, 1),
         )
diff --git a/search.py b/search.py
index 757a048..241fd43 100644
--- a/search.py
+++ b/search.py
@@ -19,10 +19,10 @@ job_plugin = studio.installed_plugins["jobs"]
 
 # Define the ranges or sets of values for each hyperparameter
 # alpha_values = list(np.round(np.linspace(2, 4, 21), 4))
-# learning_rate_values = list(np.round(np.logspace(-5, -3, 41), 5))
-alpha_values = [1]
-learning_rate_values = [5e-4]
-batch_size_values = [128]
+learning_rate_values = list(np.round(np.logspace(-5, -3, 21), 5))
+alpha_values = [0, 1, 2]
+# learning_rate_values = [5e-4]
+batch_size_values = [32, 64, 128]
 max_epochs_values = [500]
 seeds = list(range(21, 1992))