adapt grid

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"]
         )

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"))