diff --git a/main.py b/main.py
index df1420e..84915d2 100644
--- a/main.py
+++ b/main.py
@@ -40,6 +40,9 @@ def parse_args():
         default=3,
         help="Number of workers for data loading",
     )
+    parser.add_argument(
+        "--width", type=int, default=128, help="Max width of network."
+    )
 
     # Parse arguments
     args = parser.parse_args()
@@ -87,6 +90,7 @@ if __name__ == "__main__":
         alpha=args.alpha,
         learning_rate=args.lr,
         batch_size=args.bs,
+        width=args.width,
     )
 
     # Initialize model with parsed arguments
diff --git a/model.py b/model.py
index 68b9d7b..579931f 100644
--- a/model.py
+++ b/model.py
@@ -86,9 +86,9 @@ class ColorTransformerModel(pl.LightningModule):
 
         # Neural network layers
         self.network = nn.Sequential(
-            nn.Linear(3, 256),
+            nn.Linear(3, self.hparams.width),
             nn.ReLU(),
-            nn.Linear(256, 64),
+            nn.Linear(self.hparams.width, 64),
             nn.ReLU(),
             nn.Linear(64, 1),
         )
diff --git a/search.py b/search.py
index 241fd43..3936683 100644
--- a/search.py
+++ b/search.py
@@ -21,6 +21,7 @@ job_plugin = studio.installed_plugins["jobs"]
 # alpha_values = list(np.round(np.linspace(2, 4, 21), 4))
 learning_rate_values = list(np.round(np.logspace(-5, -3, 21), 5))
 alpha_values = [0, 1, 2]
+widths = [64, 128, 256, 512]
 # learning_rate_values = [5e-4]
 batch_size_values = [32, 64, 128]
 max_epochs_values = [500]
@@ -28,12 +29,13 @@ seeds = list(range(21, 1992))
 
 # Generate all possible combinations of hyperparameters
 all_params = [
-    (alpha, lr, bs, me, s)
+    (alpha, lr, bs, me, s, w)
     for alpha in alpha_values
     for lr in learning_rate_values
     for bs in batch_size_values
     for me in max_epochs_values
     for s in seeds
+    for w in widths
 ]
 
 
@@ -41,8 +43,8 @@ all_params = [
 search_params = sample(all_params, min(NUM_JOBS, len(all_params)))
 
 for idx, params in enumerate(search_params):
-    a, lr, bs, me, s = params
-    cmd = f"cd ~/colors && python main.py --alpha {a} --lr {lr} --bs {bs} --max_epochs {me} --seed {s}"
+    a, lr, bs, me, s, w = params
+    cmd = f"cd ~/colors && python main.py --alpha {a} --lr {lr} --bs {bs} --max_epochs {me} --seed {s} --width {w}"
     # job_name = f"color2_{bs}_{a}_{lr:2.2e}"
     # job_plugin.run(cmd, machine=Machine.T4, name=job_name)
     print(f"Running {params}: {cmd}")