adapt grid

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

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