chore: Update HF token and LR hyperparam schedule

neel04 · Jun 4, 2024 · 507ba6a · 507ba6a
1 parent de5c6a4
commit 507ba6a
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -1,8 +1,6 @@
 FROM python:3.11
 
 # Set environment variables
-ENV WANDB_API_KEY=78c7285b02548bf0c06dca38776c08bb6018593f
-ENV HF_TOKEN=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
 ENV jax_threefry_partitionable=1
 
 # Install base utilities

diff --git a/train_model.py b/train_model.py
@@ -136,7 +136,7 @@ def kickoff_optuna(trial, **trainer_kwargs):
     args.epochs = 1
 
     # Regularization hyperparams
-    args.lr = trial.suggest_float('lr', 1e-4, 3e-2)
+    args.lr = trial.suggest_float('lr', 1e-4, 9e-2)
     args.drop_rate = trial.suggest_float('drop_rate', 0.0, 0.1, step=0.01)
     args.weight_decay = trial.suggest_float('weight_decay', 1e-5, 1e-3)
     args.warmup_steps = trial.suggest_int('warmup_steps', 0, 500, step=100)