diff --git a/configs/grids/lp_grid.yaml b/configs/grids/lp_grid.yaml
index c5a32ec..4a35745 100644
--- a/configs/grids/lp_grid.yaml
+++ b/configs/grids/lp_grid.yaml
@@ -4,13 +4,16 @@ constants:
   logging.write_tag: linear_probe
 
 grid:
-  optimization.lr: [0.01, 0.05]
+  optimization.optimizer: [adamw]
+  optimization.lr: [0.01, 0.05, 0.001]
   optimization.weight_decay: [0.0, 5.0e-4]
-  data.batch_size: [128, 256]
+  optimization.momentum: [0.9, 0.5]
+  optimization.lr_schedule: [step, cosine]
+  data.batch_size: [256, 128]
 
 launch:
   folder: submitit_logs/
   partition: gpu1
   nodes: 1
   tasks_per_node: 1
-  time: 300
+  time: 500
diff --git a/configs/supervised_wilds_vith14_ep300-lp.yaml b/configs/supervised_wilds_vith14_ep300-lp.yaml
index 23aa965..1f1ff7d 100644
--- a/configs/supervised_wilds_vith14_ep300-lp.yaml
+++ b/configs/supervised_wilds_vith14_ep300-lp.yaml
@@ -29,7 +29,6 @@ optimization:
   epochs: 300
   lr: 0.01
   weight_decay: 5.0e-4
-  use_cosine_schedule: false
   lr_schedule: step
   step_milestones: [15, 30, 45]
   step_gamma: 0.1
diff --git a/src/train_supervised.py b/src/train_supervised.py
index f92529e..1b60ed9 100644
--- a/src/train_supervised.py
+++ b/src/train_supervised.py
@@ -249,7 +249,7 @@ def main(args, resume_preempt=False):
             milestones=o_args.get("step_milestones", [15, 30, 45]),
             gamma=o_args.get("step_gamma", 0.1),
         )
-    elif o_args["use_cosine_schedule"]:
+    elif lr_schedule == "cosine":
         scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
             optimizer, T_max=o_args["epochs"], eta_min=o_args["final_lr"]
         )
diff --git a/src/utils/logging.py b/src/utils/logging.py
index b82c198..e54d363 100644
--- a/src/utils/logging.py
+++ b/src/utils/logging.py
@@ -72,7 +72,6 @@ def build_run_name(args):
     add("wd", weight_decay)
     add("ep", epochs)
     add("sched", opt_args.get("lr_schedule"))
-    add("cos", opt_args.get("use_cosine_schedule"))
     add("ms", opt_args.get("step_milestones"))
     add("sg", opt_args.get("step_gamma"))
     add("wu", opt_args.get("warmup"))