mlp layer + new grid

configs/grids/lp_grid.yaml

@@ -4,12 +4,15 @@ constants:
   logging.write_tag: linear_probe
 
 grid:
-  optimization.optimizer: [adamw]
-  optimization.lr: [0.01, 0.05, 0.001]
-  optimization.weight_decay: [0.0, 5.0e-4]
+  optimization.optimizer: [adamw, lars, sgd]
+  optimization.lr: [0.01, 0.05]
+  optimization.weight_decay: [5.0e-4]
   optimization.momentum: [0.9, 0.5]
-  optimization.lr_schedule: [step, cosine]
-  data.batch_size: [256, 128]
+  optimization.lr_schedule: [cosine, step]
+  data.batch_size: [32, 64, 512]
+  meta.probe_type: [linear, mlp]
+  meta.mlp_hidden_dim: [1024]
+  meta.dropout: [0.0]
 
 launch:
   folder: submitit_logs/

configs/supervised_wilds_vith14_ep300-lp.yaml

@@ -6,6 +6,9 @@ meta:
   read_checkpoint: jepa-ep300.pth.tar
   use_bfloat16: true
   num_classes: 182
+  probe_type: linear
+  mlp_hidden_dim: 1024
+  dropout: 0.0
 
 data:
   batch_size: 256